|
|
|
|
|
""" |
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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}, |
|
|
|
|
|
} |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
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}" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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__ = [ |
|
|
|
|
|
"PROFESSIONAL_BACKGROUNDS", |
|
|
"create_professional_background", |
|
|
"create_gradient_background", |
|
|
|
|
|
"validate_video_file", |
|
|
"validate_video_file_detail", |
|
|
|
|
|
"segment_person_hq", |
|
|
"refine_mask_hq", |
|
|
"replace_background_hq", |
|
|
] |
|
|
|