MogensR's picture
Update utils/__init__.py
2652e08 verified
#!/usr/bin/env python3
"""
utils package (lightweight __init__)
- Export only light helpers/consts at import time
- Provide LAZY wrappers for heavy CV functions so legacy imports still work:
from utils import segment_person_hq -> OK (resolved at call time)
"""
from __future__ import annotations
import os
import logging
from typing import Dict, Any, Tuple, Optional
# import numpy as np # light; OK at import time
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------
# Background presets & builders (lightweight)
# ---------------------------------------------------------------------
PROFESSIONAL_BACKGROUNDS: Dict[str, Dict[str, Any]] = {
"office": {"color": (240, 248, 255), "gradient": True},
"studio": {"color": (32, 32, 32), "gradient": False},
"nature": {"color": (34, 139, 34), "gradient": True},
"abstract": {"color": (75, 0, 130), "gradient": True},
"white": {"color": (255, 255, 255), "gradient": False},
"black": {"color": (0, 0, 0), "gradient": False},
# add more if you like
}
def _solid_bg(color: Tuple[int,int,int], width: int, height: int) -> np.ndarray:
return np.full((height, width, 3), tuple(int(x) for x in color), dtype=np.uint8)
def _vertical_gradient(top: Tuple[int,int,int], bottom: Tuple[int,int,int], width: int, height: int) -> np.ndarray:
bg = np.zeros((height, width, 3), dtype=np.uint8)
for y in range(height):
t = y / max(1, height - 1)
r = int(top[0] * (1 - t) + bottom[0] * t)
g = int(top[1] * (1 - t) + bottom[1] * t)
b = int(top[2] * (1 - t) + bottom[2] * t)
bg[y, :] = (r, g, b)
return bg
def create_professional_background(key_or_cfg: Any, width: int, height: int) -> np.ndarray:
"""
Accepts either:
- string key in PROFESSIONAL_BACKGROUNDS
- a config dict with {"color": (r,g,b), "gradient": bool}
Returns RGB uint8 background (H, W, 3).
"""
if isinstance(key_or_cfg, str):
cfg = PROFESSIONAL_BACKGROUNDS.get(key_or_cfg, PROFESSIONAL_BACKGROUNDS["office"])
elif isinstance(key_or_cfg, dict):
cfg = key_or_cfg
else:
cfg = PROFESSIONAL_BACKGROUNDS["office"]
color = tuple(int(x) for x in cfg.get("color", (255, 255, 255)))
use_grad = bool(cfg.get("gradient", False))
if not use_grad:
return _solid_bg(color, width, height)
# simple vertical gradient dark->color
dark = (int(color[0]*0.7), int(color[1]*0.7), int(color[2]*0.7))
return _vertical_gradient(dark, color, width, height)
def create_gradient_background(spec: Dict[str, Any], width: int, height: int) -> np.ndarray:
"""
spec: {"type": "linear"|"radial", "start": (r,g,b)|"#RRGGBB", "end": (r,g,b)|"#RRGGBB", "angle_deg": float}
Returns RGB uint8 background (H, W, 3). (Radial treated as linear fallback unless extended.)
"""
import re
import cv2 # import locally to keep top-level light
def _to_rgb(c):
if isinstance(c, (list, tuple)) and len(c) == 3:
return tuple(int(x) for x in c)
if isinstance(c, str) and re.match(r"^#[0-9a-fA-F]{6}$", c):
return tuple(int(c[i:i+2], 16) for i in (1,3,5))
return (255, 255, 255)
start = _to_rgb(spec.get("start", (32, 32, 32)))
end = _to_rgb(spec.get("end", (200, 200, 200)))
angle = float(spec.get("angle_deg", 0.0))
bg = _vertical_gradient(start, end, width, height)
# rotate by angle
center = (width / 2, height / 2)
rot = cv2.getRotationMatrix2D(center, angle, 1.0)
bg = cv2.warpAffine(bg, rot, (width, height), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
return bg
# ---------------------------------------------------------------------
# Video validation (lightweight)
# ---------------------------------------------------------------------
def validate_video_file(video_path: str) -> bool:
"""
Fast sanity check: file exists, cv2 can open, first frame is readable.
Returns True/False (lightweight for UI).
"""
try:
if not video_path or not os.path.exists(video_path):
return False
import cv2 # local import
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return False
ok, frame = cap.read()
cap.release()
return bool(ok and frame is not None)
except Exception as e:
logger.warning("validate_video_file error: %s", e)
return False
def validate_video_file_detail(video_path: str) -> Tuple[bool, str]:
if not video_path:
return False, "No path provided"
if not os.path.exists(video_path):
return False, "File does not exist"
try:
import cv2
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return False, "cv2 could not open file"
ok, frame = cap.read()
cap.release()
if not ok or frame is None:
return False, "Could not read first frame"
return True, "OK"
except Exception as e:
return False, f"cv2 error: {e}"
# ---------------------------------------------------------------------
# LAZY WRAPPERS (avoid importing utils.cv_processing at module import time)
# ---------------------------------------------------------------------
def segment_person_hq(*args, **kwargs):
from .cv_processing import segment_person_hq as _f
return _f(*args, **kwargs)
def refine_mask_hq(*args, **kwargs):
from .cv_processing import refine_mask_hq as _f
return _f(*args, **kwargs)
def replace_background_hq(*args, **kwargs):
from .cv_processing import replace_background_hq as _f
return _f(*args, **kwargs)
__all__ = [
# backgrounds
"PROFESSIONAL_BACKGROUNDS",
"create_professional_background",
"create_gradient_background",
# validation
"validate_video_file",
"validate_video_file_detail",
# lazy CV exports (back-compat)
"segment_person_hq",
"refine_mask_hq",
"replace_background_hq",
]