Update utils/cv_processing.py

utils/cv_processing.py

@@ -101,6 +101,35 @@ class BackgroundReplacementError(Exception):
     """Custom exception for background replacement failures"""
     pass
 
+# ============================================================================
+# BACKGROUND HELPERS (NEW)
+# ============================================================================
+
+def _fit_image_letterbox(img_rgb: np.ndarray, dst_w: int, dst_h: int, fill=(32, 32, 32)) -> np.ndarray:
+    """
+    Fit an RGB image into (dst_h, dst_w) with letterboxing (no stretch), borders filled with `fill`.
+    """
+    h, w = img_rgb.shape[:2]
+    if h == 0 or w == 0:
+        return np.full((dst_h, dst_w, 3), fill, dtype=np.uint8)
+
+    src_aspect = w / max(1, h)
+    dst_aspect = dst_w / max(1, dst_h)
+
+    if src_aspect > dst_aspect:
+        new_w = dst_w
+        new_h = int(round(dst_w / src_aspect))
+    else:
+        new_h = dst_h
+        new_w = int(round(dst_h * src_aspect))
+
+    resized = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    canvas = np.full((dst_h, dst_w, 3), fill, dtype=np.uint8)
+    y0 = (dst_h - new_h) // 2
+    x0 = (dst_w - new_w) // 2
+    canvas[y0:y0+new_h, x0:x0+new_w] = resized
+    return canvas
+
 # ============================================================================
 # MAIN SEGMENTATION FUNCTIONS
 # ============================================================================
@@ -266,6 +295,7 @@ def segment_person_hq_original(image: np.ndarray, predictor: Any, fallback_enabl
             return _fallback_segmentation(image)
         else:
             raise SegmentationError(f"Unexpected error: {e}")
+
 # ============================================================================
 # MASK REFINEMENT FUNCTIONS
 # ============================================================================
@@ -379,11 +409,8 @@ def _matanyone_refine(image: np.ndarray, mask: np.ndarray, matanyone_processor:
         objects = [1]  # single object id
         with torch.no_grad():
             output_prob = matanyone_processor.step(img_tensor, mask_tensor, objects=objects)
-            # MatAnyOne returns output_prob as tensor
-
             refined_mask_tensor = matanyone_processor.output_prob_to_mask(output_prob)
 
-        # Convert to numpy and to uint8
         refined_mask = refined_mask_tensor.squeeze().detach().cpu().numpy()
         if refined_mask.max() <= 1.0:
             refined_mask = (refined_mask * 255).astype(np.uint8)
@@ -442,23 +469,56 @@ def replace_background_hq(frame: np.ndarray, mask: np.ndarray, background: np.nd
         else:
             raise BackgroundReplacementError(f"Unexpected error: {e}")
 
-def create_professional_background(bg_config: Dict[str, Any], width: int, height: int) -> np.ndarray:
-    """Enhanced professional background creation with quality improvements"""
-    try:
-        if bg_config["type"] == "color":
-            background = _create_solid_background(bg_config, width, height)
-        elif bg_config["type"] == "gradient":
-            background = _create_gradient_background_enhanced(bg_config, width, height)
-        else:
-            background = np.full((height, width, 3), (128, 128, 128), dtype=np.uint8)
-        
-        background = _apply_background_adjustments(background, bg_config)
-        
-        return background
-        
-    except Exception as e:
-        logger.error(f"Background creation error: {e}")
-        return np.full((height, width, 3), (128, 128, 128), dtype=np.uint8)
+def create_professional_background(bg_config: Dict[str, Any] | str, width: int, height: int) -> np.ndarray:
+    """
+    Central background builder.
+    - Accepts a style string OR a dict like:
+        {'background_choice': 'minimalist', 'custom_path': '/path/to/image.jpg'}
+      (also backwards compatible with older dicts that contained 'type'/'colors')
+    - If 'custom_path' exists, we load that image and letterbox-fit it.
+    - Returns RGB np.ndarray of shape (height, width, 3).
+    """
+    # Normalize inputs
+    choice = "minimalist"
+    custom_path = None
+
+    if isinstance(bg_config, dict):
+        # new form
+        choice = bg_config.get("background_choice", bg_config.get("name", "minimalist"))
+        custom_path = bg_config.get("custom_path")
+        # Custom image takes precedence
+        if custom_path and os.path.exists(custom_path):
+            img_bgr = cv2.imread(custom_path, cv2.IMREAD_COLOR)
+            if img_bgr is not None:
+                img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+                return _fit_image_letterbox(img_rgb, width, height, fill=(32, 32, 32))
+            else:
+                logger.warning(f"Failed to read custom background at {custom_path}. Falling back to style.")
+        # old form (has type/colors) – build from spec if present
+        if "type" in bg_config and "colors" in bg_config:
+            if bg_config["type"] == "color":
+                background = _create_solid_background(bg_config, width, height)
+            else:
+                background = _create_gradient_background_enhanced(bg_config, width, height)
+            return _apply_background_adjustments(background, bg_config)
+
+    elif isinstance(bg_config, str):
+        choice = bg_config
+
+    # No custom path → use our lightweight styles
+    choice = (choice or "minimalist").lower()
+    if choice not in PROFESSIONAL_BACKGROUNDS:
+        choice = "minimalist"
+    cfg = PROFESSIONAL_BACKGROUNDS[choice]
+
+    if cfg.get("type") == "color":
+        background = _create_solid_background(cfg, width, height)
+    else:
+        background = _create_gradient_background_enhanced(cfg, width, height)
+
+    background = _apply_background_adjustments(background, cfg)
+    return background
+
 # ============================================================================
 # VALIDATION FUNCTION
 # ============================================================================
@@ -730,6 +790,7 @@ def _auto_refine_mask_iteratively(image: np.ndarray, initial_mask: np.ndarray,
     except Exception as e:
         logger.warning(f"Iterative refinement failed: {e}")
         return initial_mask
+
 def _assess_mask_quality(mask: np.ndarray, image: np.ndarray) -> float:
     """Assess mask quality automatically"""
     try:
@@ -863,35 +924,35 @@ def _process_mask(mask: np.ndarray) -> np.ndarray:
         return fallback
 
 def _validate_mask_quality(mask: np.ndarray, image_shape: Tuple[int, int]) -> bool:
-    """Validate that the mask meets quality criteria"""
+    """
+    Soft validator: only reject clearly broken masks.
+    - Accept area ratios roughly between 2% and 95%.
+    - Don't fail on center; just warn.
+    """
     try:
         h, w = image_shape
-        mask_area = np.sum(mask > 127)
-        total_area = h * w
-        
+        total_area = max(1, h * w)
+        mask_area = int(np.sum(mask > 127))
         area_ratio = mask_area / total_area
-        if area_ratio < 0.05 or area_ratio > 0.8:
-            logger.warning(f"Suspicious mask area ratio: {area_ratio:.3f}")
-            return False
-        
-        mask_binary = mask > 127
-        mask_center_y, mask_center_x = np.where(mask_binary)
-        
-        if len(mask_center_y) == 0:
-            logger.warning("Empty mask")
+
+        # Only reject extreme cases
+        if area_ratio < 0.02 or area_ratio > 0.95:
+            logger.warning(f"Suspicious mask area ratio (hard reject): {area_ratio:.3f}")
             return False
-        
-        center_y = np.mean(mask_center_y)
-        center_x = np.mean(mask_center_x)
-        
-        if center_y < h * 0.2 or center_y > h * 0.9:
-            logger.warning(f"Mask center too far from expected person location: y={center_y/h:.2f}")
+
+        ys, xs = np.where(mask > 127)
+        if len(ys) == 0:
+            logger.warning("Empty mask (hard reject)")
             return False
-        
+
+        cy, cx = float(np.mean(ys)) / h, float(np.mean(xs)) / w
+        if cy < 0.10 or cy > 0.98:
+            logger.warning(f"Mask center unusual (advisory): y={cy:.2f}")
+
         return True
-        
+
     except Exception as e:
-        logger.warning(f"Mask validation error: {e}")
+        logger.warning(f"Mask validation error (allowing): {e}")
         return True
 
 def _fallback_segmentation(image: np.ndarray) -> np.ndarray:
@@ -1064,6 +1125,7 @@ def _simple_compositing(frame: np.ndarray, mask: np.ndarray, background: np.ndar
     except Exception as e:
         logger.error(f"Simple compositing failed: {e}")
         return frame
+
 # ============================================================================
 # HELPER FUNCTIONS - BACKGROUND CREATION
 # ============================================================================
@@ -1121,12 +1183,10 @@ def _create_vertical_gradient(colors: list, width: int, height: int) -> np.ndarr
 def _create_horizontal_gradient(colors: list, width: int, height: int) -> np.ndarray:
     """Create horizontal gradient using NumPy for performance"""
     gradient = np.zeros((height, width, 3), dtype=np.uint8)
-    
     for x in range(width):
         progress = x / width if width > 0 else 0
         color = _interpolate_color(colors, progress)
         gradient[:, x] = color
-    
     return gradient
 
 def _create_diagonal_gradient(colors: list, width: int, height: int) -> np.ndarray:
@@ -1148,7 +1208,7 @@ def _create_radial_gradient(colors: list, width: int, height: int, soft: bool =
     max_distance = np.sqrt(center_x**2 + center_y**2)
     
     y_coords, x_coords = np.mgrid[0:height, 0:width]
-    distances = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
+    distances = np.sqrt((x - center_x)**2 + (y - center_y)**2)
     progress = distances / max_distance
     progress = np.clip(progress, 0, 1)
     
@@ -1217,4 +1277,4 @@ def _apply_background_adjustments(background: np.ndarray, bg_config: Dict[str, A
         
     except Exception as e:
         logger.warning(f"Background adjustment failed: {e}")
-        return background 
+        return background