Update utilities.py

utilities.py

@@ -1,53 +1,24 @@
 #!/usr/bin/env python3
 """
 utilities.py - Helper functions and utilities for Video Background Replacement
-Contains all the utility functions, model loading, and background creation functions
-FIXED: SAM2 configuration loading issue
+Contains all the utility functions, background creation functions
+UPDATED: Models passed as parameters instead of globals
 """
 
 import os
-import sys
-import tempfile
 import cv2
 import numpy as np
-from pathlib import Path
 import torch
 import requests
-from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
-import json
-import traceback
-import time
-import shutil
-import gc
-import threading
-import queue
-from typing import Optional, Tuple, Dict, Any
+from PIL import Image, ImageDraw
 import logging
+import time
 
-# Fix OpenMP threads issue - remove problematic environment variable
-try:
-    if 'OMP_NUM_THREADS' in os.environ:
-        del os.environ['OMP_NUM_THREADS']
-except:
-    pass
-
-# Suppress warnings and optimize for quality
-import warnings
-warnings.filterwarnings("ignore")
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:1024'
-os.environ['CUDA_LAUNCH_BLOCKING'] = '0'
-
-# Setup logging for debugging
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Global variables for models (lazy loading)
-sam2_predictor = None
-matanyone_model = None
-models_loaded = False
-loading_lock = threading.Lock()
-
-# Professional background templates - Enhanced collection
+# Professional background templates
 PROFESSIONAL_BACKGROUNDS = {
     "office_modern": {
         "name": "Modern Office",
@@ -156,427 +127,13 @@
     }
 }
 
-def download_and_setup_models():
-    """ENHANCED download and setup with multiple fallback methods and lazy loading"""
-    global sam2_predictor, matanyone_model, models_loaded
-
-    with loading_lock:
-        if models_loaded:
-            return "✅ High-quality models already loaded"
-
-        try:
-            logger.info("🔄 Starting ENHANCED model loading with multiple fallbacks...")
-
-            # Check environment and system capabilities
-            is_hf_space = os.getenv("SPACE_ID") is not None
-            is_colab = 'google.colab' in sys.modules
-            is_kaggle = os.environ.get('KAGGLE_KERNEL_RUN_TYPE') is not None
-
-            env_type = "HuggingFace Space" if is_hf_space else "Google Colab" if is_colab else "Kaggle" if is_kaggle else "Local"
-            logger.info(f"Environment detected: {env_type}")
-
-            # Load PyTorch and check GPU
-            import torch
-            logger.info(f"✅ PyTorch {torch.__version__} - CUDA: {torch.cuda.is_available()}")
-
-            if torch.cuda.is_available():
-                try:
-                    gpu_name = torch.cuda.get_device_name(0)
-                    gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
-                    logger.info(f"🎮 GPU: {gpu_name} ({gpu_memory:.1f}GB)")
-                except Exception as e:
-                    logger.info(f"🎮 GPU available but details unavailable: {e}")
-
-            # === ENHANCED SAM2 LOADING WITH MULTIPLE METHODS ===
-            sam2_loaded = False
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-
-            # --- Improved YAML/config handling ---
-            config_paths = [
-                "./configs",            # Local ./configs directory
-                "/home/user/app/configs",  # Typical in HF spaces
-                os.path.expanduser("~/.cache/sam2/configs"),
-            ]
-            config_dir = None
-            for path in config_paths:
-                if os.path.isdir(path):
-                    config_dir = path
-                    break
-
-            # Copy bundled .yaml files to a found config_dir if not present
-            bundled_configs = ["sam2_hiera_large.yaml", "sam2_hiera_tiny.yaml"]
-            if config_dir:
-                for cfg_file in bundled_configs:
-                    src = Path(cfg_file)
-                    dest = Path(config_dir) / cfg_file
-                    if src.exists() and not dest.exists():
-                        shutil.copyfile(src, dest)
-                        logger.info(f"✅ Copied {cfg_file} to {config_dir}")
-            else:
-                logger.warning("No configs directory found for SAM2! Fallback to default logic.")
-
-            # --- Method 1: Try direct import (requirements.txt installation) ---
-            try:
-                logger.info("🔄 SAM2 Method 1: Direct import from requirements...")
-                from sam2.build_sam import build_sam2
-                from sam2.sam2_image_predictor import SAM2ImagePredictor
-                sam2_loaded = True
-                logger.info("✅ SAM2 imported directly from installed package")
-            except ImportError as e:
-                logger.info(f"❌ SAM2 Method 1 failed: {e}")
-
-            # --- Method 2: Clone and properly setup SAM2 ---
-            if not sam2_loaded:
-                try:
-                    logger.info("🔄 SAM2 Method 2: Cloning and setting up repository...")
-                    sam2_dir = "/tmp/segment-anything-2"
-                    if not os.path.exists(sam2_dir):
-                        logger.info("📥 Cloning SAM2 repository...")
-                        clone_cmd = f"git clone --depth 1 https://github.com/facebookresearch/segment-anything-2.git {sam2_dir}"
-                        result = os.system(clone_cmd)
-                        if result == 0:
-                            logger.info("✅ SAM2 repository cloned successfully")
-                        else:
-                            raise Exception("Git clone failed")
-                    # Add to path
-                    if sam2_dir not in sys.path:
-                        sys.path.insert(0, sam2_dir)
-                    # Install SAM2 dependencies if needed
-                    try:
-                        import hydra
-                    except ImportError:
-                        logger.info("Installing Hydra-core for SAM2 configs...")
-                        os.system("pip install hydra-core --quiet")
-                    from sam2.build_sam import build_sam2
-                    from sam2.sam2_image_predictor import SAM2ImagePredictor
-                    sam2_loaded = True
-                    logger.info("✅ SAM2 imported after cloning")
-                except Exception as e:
-                    logger.info(f"❌ SAM2 Method 2 failed: {e}")
-
-            # --- Method 3: Use simplified SAM2 loading without Hydra configs ---
-            if not sam2_loaded:
-                try:
-                    logger.info("🔄 SAM2 Method 3: Simplified loading without Hydra...")
-                    # Download checkpoint first
-                    cache_dir = os.path.expanduser("~/.cache/sam2")
-                    os.makedirs(cache_dir, exist_ok=True)
-                    checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
-                    sam2_checkpoint = os.path.join(cache_dir, "sam2_hiera_tiny.pt")
-                    if not os.path.exists(sam2_checkpoint):
-                        logger.info("📥 Downloading SAM2 checkpoint...")
-                        response = requests.get(checkpoint_url, stream=True)
-                        with open(sam2_checkpoint, 'wb') as f:
-                            for chunk in response.iter_content(chunk_size=8192):
-                                if chunk:
-                                    f.write(chunk)
-                        logger.info("✅ Checkpoint downloaded")
-                    checkpoint = torch.load(sam2_checkpoint, map_location=device)
-                    class SimpleSAM2Predictor:
-                        def __init__(self, checkpoint_path, device):
-                            self.device = device
-                            self.checkpoint_path = checkpoint_path
-                            self.image = None
-                            logger.info("Using simplified SAM2 predictor")
-                        def set_image(self, image):
-                            self.image = image.copy()
-                        def predict(self, point_coords, point_labels, multimask_output=True):
-                            if self.image is None:
-                                raise ValueError("No image set")
-                            h, w = self.image.shape[:2]
-                            mask = np.zeros((h, w), dtype=np.uint8)
-                            for point in point_coords:
-                                x, y = int(point[0]), int(point[1])
-                                cv2.circle(mask, (x, y), min(w, h)//4, 255, -1)
-                            try:
-                                mask_3d = np.zeros((h, w), dtype=np.uint8)
-                                mask_3d[mask > 0] = cv2.GC_PR_FGD
-                                mask_3d[mask == 0] = cv2.GC_PR_BGD
-                                bgd_model = np.zeros((1, 65), np.float64)
-                                fgd_model = np.zeros((1, 65), np.float64)
-                                cv2.grabCut(self.image, mask_3d, None, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_MASK)
-                                mask = np.where((mask_3d == cv2.GC_FGD) | (mask_3d == cv2.GC_PR_FGD), 255, 0).astype('uint8')
-                            except:
-                                pass
-                            return [mask], [1.0], None
-                    sam2_predictor = SimpleSAM2Predictor(sam2_checkpoint, device)
-                    sam2_loaded = True
-                    logger.info("✅ Using simplified SAM2 predictor")
-                except Exception as e:
-                    logger.info(f"❌ SAM2 Method 3 failed: {e}")
-
-            # --- Method 4: Install via pip and try again ---
-            if not sam2_loaded:
-                try:
-                    logger.info("🔄 SAM2 Method 4: Installing via pip...")
-                    os.system("pip install hydra-core omegaconf --quiet")
-                    sam2_dir = "/tmp/sam2_install"
-                    if os.path.exists(sam2_dir):
-                        shutil.rmtree(sam2_dir)
-                    clone_cmd = f"git clone https://github.com/facebookresearch/segment-anything-2.git {sam2_dir}"
-                    os.system(clone_cmd)
-                    install_cmd = f"cd {sam2_dir} && pip install -e . --quiet"
-                    os.system(install_cmd)
-                    from sam2.build_sam import build_sam2
-                    from sam2.sam2_image_predictor import SAM2ImagePredictor
-                    sam2_loaded = True
-                    logger.info("✅ SAM2 installed and imported via pip")
-                except Exception as e:
-                    logger.info(f"❌ SAM2 Method 4 failed: {e}")
-
-            if not sam2_loaded:
-                logger.warning("❌ All SAM2 loading methods failed, using OpenCV fallback")
-                sam2_predictor = create_opencv_segmentation_fallback()
-            else:
-                if not isinstance(sam2_predictor, object) or sam2_predictor is None:
-                    try:
-                        # Choose model size based on environment and resources
-                        if (is_hf_space and not torch.cuda.is_available()) or device == "cpu":
-                            model_name = "sam2_hiera_tiny"
-                            checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
-                            logger.info("🔧 Using SAM2 Tiny for CPU/limited resources")
-                        else:
-                            model_name = "sam2_hiera_large"
-                            checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"
-                            logger.info("🔧 Using SAM2 Large for maximum quality")
-                        # Download checkpoint
-                        cache_dir = os.path.expanduser("~/.cache/sam2")
-                        os.makedirs(cache_dir, exist_ok=True)
-                        sam2_checkpoint = os.path.join(cache_dir, f"{model_name}.pt")
-                        if not os.path.exists(sam2_checkpoint):
-                            logger.info(f"📥 Downloading {model_name} checkpoint...")
-                            try:
-                                response = requests.get(checkpoint_url, stream=True)
-                                total_size = int(response.headers.get('content-length', 0))
-                                downloaded = 0
-                                with open(sam2_checkpoint, 'wb') as f:
-                                    for chunk in response.iter_content(chunk_size=8192):
-                                        if chunk:
-                                            f.write(chunk)
-                                            downloaded += len(chunk)
-                                            if total_size > 0 and downloaded % (total_size // 20) < 8192:
-                                                percent = (downloaded / total_size) * 100
-                                                logger.info(f"📥 Download progress: {percent:.1f}%")
-                                logger.info(f"✅ {model_name} downloaded successfully")
-                            except Exception as e:
-                                logger.error(f"❌ Download failed: {e}")
-                                raise
-                        else:
-                            logger.info(f"✅ Using cached {model_name}")
-                        # Load SAM2 model - use the config name without .yaml extension
-                        try:
-                            logger.info(f"🚀 Loading SAM2 {model_name} on {device}...")
-                            model_cfg = model_name
-                            if device == "cpu" or is_hf_space:
-                                torch.set_num_threads(min(4, os.cpu_count() or 1))
-                                if torch.cuda.is_available():
-                                    torch.cuda.empty_cache()
-                            sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)
-                            sam2_predictor = SAM2ImagePredictor(sam2_model)
-                            logger.info(f"✅ SAM2 model loaded successfully on {device}")
-                        except Exception as e:
-                            if device == "cuda":
-                                logger.warning(f"❌ GPU loading failed: {e}")
-                                logger.info("🔄 Trying CPU fallback...")
-                                try:
-                                    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
-                                    sam2_predictor = SAM2ImagePredictor(sam2_model)
-                                    device = "cpu"
-                                    logger.info("✅ SAM2 loaded on CPU fallback")
-                                except Exception as e2:
-                                    logger.error(f"❌ CPU fallback also failed: {e2}")
-                                    logger.info("🔄 Using OpenCV segmentation fallback")
-                                    sam2_predictor = create_opencv_segmentation_fallback()
-                            else:
-                                logger.error(f"❌ SAM2 loading failed: {e}")
-                                logger.info("🔄 Using OpenCV segmentation fallback")
-                                sam2_predictor = create_opencv_segmentation_fallback()
-                    except Exception as e:
-                        logger.error(f"❌ SAM2 initialization failed: {e}")
-                        sam2_predictor = create_opencv_segmentation_fallback()
-
-            # === ENHANCED MATANYONE LOADING WITH MULTIPLE METHODS ===
-            matanyone_loaded = False
-            # Method 1: Try HuggingFace Hub integration
-            try:
-                logger.info("🔄 MatAnyone Method 1: HuggingFace Hub...")
-                from huggingface_hub import hf_hub_download
-                from matanyone import InferenceCore
-                matanyone_model = InferenceCore("PeiqingYang/MatAnyone")
-                matanyone_loaded = True
-                logger.info("✅ MatAnyone loaded via HuggingFace Hub")
-            except Exception as e:
-                logger.info(f"❌ MatAnyone Method 1 failed: {e}")
-            # Method 2: Try direct import
-            if not matanyone_loaded:
-                try:
-                    logger.info("🔄 MatAnyone Method 2: Direct import...")
-                    matanyone_paths = [
-                        '/tmp/MatAnyone',
-                        './MatAnyone',
-                        '/content/MatAnyone',
-                        '/kaggle/working/MatAnyone'
-                    ]
-                    for path in matanyone_paths:
-                        if os.path.exists(path):
-                            sys.path.append(path)
-                            break
-                    from inference import MatAnyoneInference
-                    matanyone_model = MatAnyoneInference()
-                    matanyone_loaded = True
-                    logger.info("✅ MatAnyone loaded via direct import")
-                except Exception as e:
-                    logger.info(f"❌ MatAnyone Method 2 failed: {e}")
-            # Method 3: Try GitHub installation
-            if not matanyone_loaded:
-                try:
-                    logger.info("🔄 MatAnyone Method 3: Installing from GitHub...")
-                    install_cmd = "pip install git+https://github.com/pq-yang/MatAnyone.git"
-                    result = os.system(install_cmd)
-                    if result == 0:
-                        from matanyone import InferenceCore
-                        matanyone_model = InferenceCore("PeiqingYang/MatAnyone")
-                        matanyone_loaded = True
-                        logger.info("✅ MatAnyone installed and loaded via GitHub")
-                    else:
-                        raise Exception("GitHub install failed")
-                except Exception as e:
-                    logger.info(f"❌ MatAnyone Method 3 failed: {e}")
-            # Method 4: Enhanced OpenCV fallback (CINEMA QUALITY)
-            if not matanyone_loaded:
-                logger.info("🎨 Using ENHANCED OpenCV fallback for cinema-quality matting...")
-                matanyone_model = create_enhanced_matting_fallback()
-                matanyone_loaded = True
-
-            # Memory cleanup
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            models_loaded = True
-            gpu_info = ""
-            if torch.cuda.is_available() and device == "cuda":
-                try:
-                    gpu_info = f" (GPU: {torch.cuda.get_device_name(0)})"
-                except:
-                    gpu_info = " (GPU)"
-            else:
-                gpu_info = " (CPU)"
-            success_msg = f"✅ ENHANCED high-quality models loaded successfully!{gpu_info}"
-            logger.info(success_msg)
-            return success_msg
-
-        except Exception as e:
-            error_msg = f"❌ Enhanced loading failed: {str(e)}"
-            logger.error(error_msg)
-            logger.error(f"Full traceback: {traceback.format_exc()}")
-            return error_msg
-
-# ... next: create_opencv_segmentation_fallback(), create_enhanced_matting_fallback(), and much more
-
-def create_opencv_segmentation_fallback():
-    """Create comprehensive OpenCV-based segmentation fallback"""
-    class OpenCVSegmentationFallback:
-        def __init__(self):
-            logger.info("🔧 Initializing OpenCV segmentation fallback")
-            self.bg_subtractor = cv2.createBackgroundSubtractorMOG2(detectShadows=True)
-            self.image = None
-
-        def set_image(self, image):
-            self.image = image.copy()
-
-        def predict(self, point_coords, point_labels, multimask_output=True):
-            """Advanced OpenCV-based person segmentation with multiple techniques"""
-            if self.image is None:
-                raise ValueError("No image set")
-            h, w = self.image.shape[:2]
-            try:
-                # Multi-method segmentation
-                masks = []
-                # Skin tone detection (HSV ranges)
-                hsv = cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV)
-                lower_skin1 = np.array([0, 20, 70], dtype=np.uint8)
-                upper_skin1 = np.array([20, 255, 255], dtype=np.uint8)
-                lower_skin2 = np.array([0, 20, 70], dtype=np.uint8)
-                upper_skin2 = np.array([25, 255, 255], dtype=np.uint8)
-                skin_mask1 = cv2.inRange(hsv, lower_skin1, upper_skin1)
-                skin_mask2 = cv2.inRange(hsv, lower_skin2, upper_skin2)
-                skin_mask = cv2.bitwise_or(skin_mask1, skin_mask2)
-                # Edge detection for person outline
-                gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
-                edges = cv2.Canny(gray, 50, 150)
-                # Focus on center region (with point guidance)
-                center_x, center_y = w//2, h//2
-                if len(point_coords) > 0:
-                    center_x = int(np.mean(point_coords[:, 0]))
-                    center_y = int(np.mean(point_coords[:, 1]))
-                center_mask = np.zeros((h, w), dtype=np.uint8)
-                roi_width = w // 3
-                roi_height = h // 2
-                cv2.ellipse(center_mask, (center_x, center_y), (roi_width, roi_height), 0, 0, 360, 255, -1)
-                # Combine masks
-                combined_mask = cv2.bitwise_or(skin_mask, edges // 4)
-                combined_mask = cv2.bitwise_and(combined_mask, center_mask)
-                kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
-                kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-                combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel_close)
-                combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel_open)
-                contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-                if contours:
-                    largest_contour = max(contours, key=cv2.contourArea)
-                    mask = np.zeros((h, w), dtype=np.uint8)
-                    cv2.fillPoly(mask, [largest_contour], 255)
-                    mask = cv2.GaussianBlur(mask, (5, 5), 2.0)
-                    mask = (mask > 127).astype(np.uint8)
-                else:
-                    mask = center_mask
-                mask = cv2.medianBlur(mask, 5)
-                masks.append(mask)
-                scores = [1.0]
-                return masks, scores, None
-            except Exception as e:
-                logger.warning(f"OpenCV segmentation error: {e}")
-                mask = np.zeros((h, w), dtype=np.uint8)
-                x1, y1 = w//4, h//6
-                x2, y2 = 3*w//4, 5*h//6
-                mask[y1:y2, x1:x2] = 255
-                return [mask], [1.0], None
-    return OpenCVSegmentationFallback()
-
-def create_enhanced_matting_fallback():
-    """Create enhanced matting fallback with advanced OpenCV techniques"""
-    class EnhancedMattingFallback:
-        def __init__(self):
-            logger.info("🎨 Initializing enhanced matting fallback")
-        def infer(self, image, mask):
-            """Enhanced mask refinement using advanced OpenCV techniques"""
-            try:
-                if len(mask.shape) == 3:
-                    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-                refined_mask = cv2.bilateralFilter(mask, 9, 75, 75)
-                kernel_ellipse = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
-                refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel_ellipse)
-                refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_OPEN, kernel_ellipse)
-                refined_mask = cv2.GaussianBlur(refined_mask, (3, 3), 1.0)
-                edges = cv2.Canny(refined_mask, 50, 150)
-                edge_enhancement = cv2.dilate(edges, np.ones((2, 2), np.uint8), iterations=1)
-                refined_mask = cv2.bitwise_or(refined_mask, edge_enhancement // 4)
-                dist_transform = cv2.distanceTransform(refined_mask, cv2.DIST_L2, 5)
-                dist_transform = cv2.normalize(dist_transform, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
-                alpha = 0.7
-                refined_mask = cv2.addWeighted(refined_mask, alpha, dist_transform, 1-alpha, 0)
-                refined_mask = cv2.medianBlur(refined_mask, 3)
-                refined_mask = cv2.GaussianBlur(refined_mask, (1, 1), 0.5)
-                return refined_mask
-            except Exception as e:
-                logger.warning(f"Enhanced matting error: {e}")
-                return mask if len(mask.shape) == 2 else cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-    return EnhancedMattingFallback()
-
-def segment_person_hq(image):
-    """High-quality person segmentation using SAM2 or fallback with optimized points"""
+def segment_person_hq(image, predictor):
+    """High-quality person segmentation using provided SAM2 predictor"""
     try:
-        sam2_predictor.set_image(image)
+        predictor.set_image(image)
         h, w = image.shape[:2]
+        
+        # Strategic point placement for person detection
         points = np.array([
             [w//2, h//4],    # Top-center (head)
             [w//2, h//2],    # Center (torso)
@@ -587,23 +144,33 @@ def segment_person_hq(image):
             [4*w//5, h//5]   # Top-right (hair/accessories)
         ])
         labels = np.ones(len(points))
-        masks, scores, _ = sam2_predictor.predict(
+        
+        masks, scores, _ = predictor.predict(
             point_coords=points,
             point_labels=labels,
             multimask_output=True
         )
+        
+        # Select best mask
         best_idx = np.argmax(scores)
         best_mask = masks[best_idx]
+        
+        # Ensure proper format
         if len(best_mask.shape) > 2:
             best_mask = best_mask.squeeze()
         if best_mask.dtype != np.uint8:
             best_mask = (best_mask * 255).astype(np.uint8)
+        
+        # Post-process mask
         kernel = np.ones((3, 3), np.uint8)
         best_mask = cv2.morphologyEx(best_mask, cv2.MORPH_CLOSE, kernel)
         best_mask = cv2.GaussianBlur(best_mask.astype(np.float32), (3, 3), 0.8)
+        
         return (best_mask * 255).astype(np.uint8) if best_mask.max() <= 1.0 else best_mask.astype(np.uint8)
+        
     except Exception as e:
         logger.error(f"Segmentation error: {e}")
+        # Fallback to simple center mask
         h, w = image.shape[:2]
         fallback_mask = np.zeros((h, w), dtype=np.uint8)
         x1, y1 = w//4, h//6
@@ -611,26 +178,129 @@ def segment_person_hq(image):
         fallback_mask[y1:y2, x1:x2] = 255
         return fallback_mask
 
-def refine_mask_hq(image, mask):
-    """Cinema-quality mask refinement with stronger edge preservation"""
+def refine_mask_hq(image, mask, matanyone_processor):
+    """Cinema-quality mask refinement using provided MatAnyone processor"""
     try:
+        # Prepare image for matting
         image_filtered = cv2.bilateralFilter(image, 10, 75, 75)
-        refined_mask = matanyone_model.infer(image_filtered, mask)
+        
+        # Use MatAnyone for refinement
+        if hasattr(matanyone_processor, 'process_video'):
+            # If it's the HF InferenceCore, we need to handle differently
+            # For now, use enhanced OpenCV refinement
+            refined_mask = enhance_mask_opencv(image_filtered, mask)
+        else:
+            # Direct inference call
+            refined_mask = matanyone_processor.infer(image_filtered, mask)
+        
+        # Ensure proper format
         if len(refined_mask.shape) == 3:
             refined_mask = cv2.cvtColor(refined_mask, cv2.COLOR_BGR2GRAY)
+        
+        # Additional refinement
         refined_mask = cv2.bilateralFilter(refined_mask, 10, 75, 75)
         refined_mask = cv2.medianBlur(refined_mask, 3)
+        
         return refined_mask
+        
     except Exception as e:
         logger.error(f"Mask refinement error: {e}")
+        return enhance_mask_opencv(image, mask)
+
+def enhance_mask_opencv(image, mask):
+    """Enhanced mask refinement using OpenCV techniques"""
+    try:
+        if len(mask.shape) == 3:
+            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+        
+        # Bilateral filtering for edge preservation
+        refined_mask = cv2.bilateralFilter(mask, 9, 75, 75)
+        
+        # Morphological operations
+        kernel_ellipse = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel_ellipse)
+        refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_OPEN, kernel_ellipse)
+        
+        # Gaussian blur for smoothing
+        refined_mask = cv2.GaussianBlur(refined_mask, (3, 3), 1.0)
+        
+        # Edge enhancement
+        edges = cv2.Canny(refined_mask, 50, 150)
+        edge_enhancement = cv2.dilate(edges, np.ones((2, 2), np.uint8), iterations=1)
+        refined_mask = cv2.bitwise_or(refined_mask, edge_enhancement // 4)
+        
+        # Distance transform for better interior
+        dist_transform = cv2.distanceTransform(refined_mask, cv2.DIST_L2, 5)
+        dist_transform = cv2.normalize(dist_transform, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+        
+        # Blend with distance transform
+        alpha = 0.7
+        refined_mask = cv2.addWeighted(refined_mask, alpha, dist_transform, 1-alpha, 0)
+        
+        # Final smoothing
+        refined_mask = cv2.medianBlur(refined_mask, 3)
+        refined_mask = cv2.GaussianBlur(refined_mask, (1, 1), 0.5)
+        
+        return refined_mask
+        
+    except Exception as e:
+        logger.warning(f"Enhanced mask refinement error: {e}")
         return mask if len(mask.shape) == 2 else cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
 
 def create_green_screen_background(frame):
-    """Create green screen background (Stage 1 of two-stage process)"""
+    """Create green screen background for two-stage processing"""
     h, w = frame.shape[:2]
     green_screen = np.full((h, w, 3), (0, 177, 64), dtype=np.uint8)
     return green_screen
 
+def replace_background_hq(frame, mask, background):
+    """High-quality background replacement with advanced compositing"""
+    try:
+        # Resize background to match frame
+        background = cv2.resize(background, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LANCZOS4)
+        
+        # Ensure mask is single channel
+        if len(mask.shape) == 3:
+            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+        
+        # Normalize mask to 0-1 range
+        mask_float = mask.astype(np.float32) / 255.0
+        
+        # Edge feathering for smooth transitions
+        feather_radius = 3
+        kernel_size = feather_radius * 2 + 1
+        mask_feathered = cv2.GaussianBlur(mask_float, (kernel_size, kernel_size), feather_radius/3)
+        
+        # Create 3-channel mask
+        mask_3channel = np.stack([mask_feathered] * 3, axis=2)
+        
+        # Linear gamma correction for proper compositing
+        frame_linear = np.power(frame.astype(np.float32) / 255.0, 2.2)
+        background_linear = np.power(background.astype(np.float32) / 255.0, 2.2)
+        
+        # Composite in linear space
+        result_linear = frame_linear * mask_3channel + background_linear * (1 - mask_3channel)
+        
+        # Convert back to gamma space
+        result = np.power(result_linear, 1/2.2) * 255.0
+        result = np.clip(result, 0, 255).astype(np.uint8)
+        
+        return result
+        
+    except Exception as e:
+        logger.error(f"Background replacement error: {e}")
+        # Fallback to simple replacement
+        try:
+            if len(mask.shape) == 3:
+                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+            background = cv2.resize(background, (frame.shape[1], frame.shape[0]))
+            mask_normalized = mask.astype(np.float32) / 255.0
+            mask_3channel = np.stack([mask_normalized] * 3, axis=2)
+            result = frame * mask_3channel + background * (1 - mask_3channel)
+            return result.astype(np.uint8)
+        except:
+            return frame
+
 def create_professional_background(bg_config, width, height):
     """Create professional background based on configuration"""
     try:
@@ -649,10 +319,11 @@ def create_professional_background(bg_config, width, height):
         return np.full((height, width, 3), (128, 128, 128), dtype=np.uint8)
 
 def create_gradient_background(bg_config, width, height):
-    """Create high-quality gradient backgrounds with comprehensive direction support"""
+    """Create high-quality gradient backgrounds"""
     try:
         colors = bg_config["colors"]
         direction = bg_config.get("direction", "vertical")
+        
         # Convert hex to RGB
         rgb_colors = []
         for color_hex in colors:
@@ -662,10 +333,14 @@ def create_gradient_background(bg_config, width, height):
                 rgb_colors.append(rgb)
             except ValueError:
                 rgb_colors.append((128, 128, 128))
+        
         if not rgb_colors:
             rgb_colors = [(128, 128, 128)]
+        
+        # Create PIL image for gradient
         pil_img = Image.new('RGB', (width, height))
         draw = ImageDraw.Draw(pil_img)
+        
         def interpolate_color(colors, progress):
             if len(colors) == 1:
                 return colors[0]
@@ -686,6 +361,8 @@ def interpolate_color(colors, progress):
                     g = int(c1[1] + (c2[1] - c1[1]) * local_progress)
                     b = int(c1[2] + (c2[2] - c1[2]) * local_progress)
                     return (r, g, b)
+        
+        # Generate gradient based on direction
         if direction == "vertical":
             for y in range(height):
                 progress = y / height if height > 0 else 0
@@ -717,165 +394,31 @@ def interpolate_color(colors, progress):
                     color = interpolate_color(rgb_colors, progress)
                     pil_img.putpixel((x, y), color)
         else:
+            # Default to vertical
             for y in range(height):
                 progress = y / height if height > 0 else 0
                 color = interpolate_color(rgb_colors, progress)
                 draw.line([(0, y), (width, y)], fill=color)
+        
+        # Convert to OpenCV format
         background = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
         return background
+        
     except Exception as e:
         logger.error(f"Gradient creation error: {e}")
+        # Fallback gradient
         background = np.zeros((height, width, 3), dtype=np.uint8)
         for y in range(height):
             intensity = int(255 * (y / height)) if height > 0 else 128
             background[y, :] = [intensity, intensity, intensity]
         return background
 
-def replace_background_hq(frame, mask, background):
-    """High-quality background replacement with advanced compositing"""
-    try:
-        background = cv2.resize(background, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LANCZOS4)
-        if len(mask.shape) == 3:
-            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-        mask_float = mask.astype(np.float32) / 255.0
-        feather_radius = 3
-        kernel_size = feather_radius * 2 + 1
-        mask_feathered = cv2.GaussianBlur(mask_float, (kernel_size, kernel_size), feather_radius/3)
-        mask_3channel = np.stack([mask_feathered] * 3, axis=2)
-        frame_linear = np.power(frame.astype(np.float32) / 255.0, 2.2)
-        background_linear = np.power(background.astype(np.float32) / 255.0, 2.2)
-        result_linear = frame_linear * mask_3channel + background_linear * (1 - mask_3channel)
-        result = np.power(result_linear, 1/2.2) * 255.0
-        result = np.clip(result, 0, 255).astype(np.uint8)
-        return result
-    except Exception as e:
-        logger.error(f"Background replacement error: {e}")
-        try:
-            if len(mask.shape) == 3:
-                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-            background = cv2.resize(background, (frame.shape[1], frame.shape[0]))
-            mask_normalized = mask.astype(np.float32) / 255.0
-            mask_3channel = np.stack([mask_normalized] * 3, axis=2)
-            result = frame * mask_3channel + background * (1 - mask_3channel)
-            return result.astype(np.uint8)
-        except:
-            return frame
-
-def get_model_status():
-    """Get current model loading status with detailed information"""
-    if models_loaded:
-        try:
-            gpu_info = ""
-            if torch.cuda.is_available():
-                try:
-                    gpu_name = torch.cuda.get_device_name(0)
-                    gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
-                    gpu_info = f" (GPU: {gpu_name[:20]}{'...' if len(gpu_name) > 20 else ''} - {gpu_memory:.1f}GB)"
-                except:
-                    gpu_info = " (GPU Available)"
-            else:
-                gpu_info = " (CPU Mode)"
-            return f"✅ ENHANCED high-quality models loaded{gpu_info}"
-        except:
-            return "✅ ENHANCED high-quality models loaded"
-    else:
-        return "⏳ Models not loaded. Click 'Load Models' for ENHANCED cinema-quality processing."
-
-def validate_video_file(video_path):
-    """Validate video file format and basic properties"""
-    if not video_path or not os.path.exists(video_path):
-        return False, "Video file not found"
-    try:
-        cap = cv2.VideoCapture(video_path)
-        if not cap.isOpened():
-            return False, "Cannot open video file"
-        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        if frame_count == 0:
-            return False, "Video appears to be empty"
-        cap.release()
-        return True, "Video file valid"
-    except Exception as e:
-        return False, f"Error validating video: {str(e)}"
-
-def cleanup_temp_files():
-    """Clean up temporary files to free disk space"""
-    try:
-        temp_patterns = [
-            "/tmp/processed_video_*.mp4",
-            "/tmp/final_output_*.mp4",
-            "/tmp/greenscreen_*.mp4",
-            "/tmp/gradient_*.png",
-            "/tmp/*.pt",  # Model checkpoints
-        ]
-        import glob
-        for pattern in temp_patterns:
-            for file_path in glob.glob(pattern):
-                try:
-                    if os.path.exists(file_path):
-                        if time.time() - os.path.getmtime(file_path) > 3600:
-                            os.remove(file_path)
-                            logger.info(f"Cleaned up: {file_path}")
-                except Exception as e:
-                    logger.warning(f"Could not clean up {file_path}: {e}")
-    except Exception as e:
-        logger.warning(f"Cleanup error: {e}")
-
-def get_available_backgrounds():
-    """Get list of available professional backgrounds"""
-    return {key: config["name"] for key, config in PROFESSIONAL_BACKGROUNDS.items()}
-
-def create_custom_gradient(colors, direction="vertical", width=1920, height=1080):
-    """
-    Create a custom gradient background
-    Args:
-        colors: List of hex colors (e.g., ["#ff0000", "#00ff00"])
-        direction: "vertical", "horizontal", "diagonal", "radial"
-        width, height: Dimensions
-    Returns:
-        numpy array of the generated background
-    """
-    try:
-        bg_config = {
-            "type": "gradient",
-            "colors": colors,
-            "direction": direction
-        }
-        return create_gradient_background(bg_config, width, height)
-    except Exception as e:
-        logger.error(f"Error creating custom gradient: {e}")
-        return None
-
-def create_directories():
-    """Create necessary directories for the application"""
-    try:
-        directories = [
-            "/tmp/MyAvatar",
-            "/tmp/MyAvatar/My_Videos",
-            os.path.expanduser("~/.cache/sam2"),
-        ]
-        for directory in directories:
-            os.makedirs(directory, exist_ok=True)
-            logger.info(f"📁 Created/verified directory: {directory}")
-        return True
-    except Exception as e:
-        logger.error(f"Error creating directories: {e}")
-        return False
-
-def optimize_memory_usage():
-    """Optimize memory usage by cleaning up unused resources"""
-    try:
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
-        cv2.ocl.setUseOpenCL(False)
-        return True
-    except Exception as e:
-        logger.warning(f"Memory optimization failed: {e}")
-        return False
 def create_procedural_background(prompt, style, width, height):
     """Create procedural background based on text prompt and style"""
     try:
         prompt_lower = prompt.lower()
+        
+        # Color mapping based on keywords
         color_map = {
             'blue': ['#1e3c72', '#2a5298', '#3498db'],
             'ocean': ['#74b9ff', '#0984e3', '#00cec9'], 
@@ -899,12 +442,15 @@ def create_procedural_background(prompt, style, width, height):
             'minimal': ['#ffffff', '#f1f2f6', '#ddd'],
             'abstract': ['#6c5ce7', '#a29bfe', '#fd79a8']
         }
-        selected_colors = ['#3498db', '#2ecc71', '#e74c3c']
+        
+        # Select colors based on prompt
+        selected_colors = ['#3498db', '#2ecc71', '#e74c3c']  # Default
         for keyword, colors in color_map.items():
             if keyword in prompt_lower:
                 selected_colors = colors
                 break
-
+        
+        # Create background based on style
         if style == "abstract":
             return create_abstract_background(selected_colors, width, height)
         elif style == "minimalist":
@@ -916,12 +462,14 @@ def create_procedural_background(prompt, style, width, height):
         elif style == "artistic":
             return create_artistic_background(selected_colors, width, height)
         else:
+            # Default gradient
             bg_config = {
                 "type": "gradient",
                 "colors": selected_colors[:2],
                 "direction": "diagonal"
             }
             return create_gradient_background(bg_config, width, height)
+            
     except Exception as e:
         logger.error(f"Procedural background creation failed: {e}")
         return None
@@ -930,12 +478,16 @@ def create_abstract_background(colors, width, height):
     """Create abstract geometric background"""
     try:
         background = np.zeros((height, width, 3), dtype=np.uint8)
+        
+        # Convert hex colors to BGR
         bgr_colors = []
         for color in colors:
             hex_color = color.lstrip('#')
             rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
             bgr = rgb[::-1]
             bgr_colors.append(bgr)
+        
+        # Create base gradient
         for y in range(height):
             progress = y / height
             color = [
@@ -943,17 +495,22 @@ def create_abstract_background(colors, width, height):
                 for i in range(3)
             ]
             background[y, :] = color
+        
+        # Add geometric shapes
         import random
-        random.seed(42)
+        random.seed(42)  # Consistent results
         for _ in range(8):
             center_x = random.randint(width//4, 3*width//4)
             center_y = random.randint(height//4, 3*height//4)
             radius = random.randint(width//20, width//8)
             color = bgr_colors[random.randint(0, len(bgr_colors)-1)]
+            
             overlay = background.copy()
             cv2.circle(overlay, (center_x, center_y), radius, color, -1)
             cv2.addWeighted(background, 0.7, overlay, 0.3, 0, background)
+        
         return background
+        
     except Exception as e:
         logger.error(f"Abstract background creation failed: {e}")
         return None
@@ -966,15 +523,8 @@ def create_minimalist_background(colors, width, height):
             "colors": colors[:2],
             "direction": "soft_radial"
         }
-        background = create_gradient_background(bg_config, width, height)
-        overlay = background.copy()
-        center_x, center_y = width//2, height//2
-        hex_color = colors[0].lstrip('#')
-        rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
-        bgr = rgb[::-1]
-        cv2.circle(overlay, (center_x, center_y), min(width, height)//3, bgr, -1)
-        cv2.addWeighted(background, 0.95, overlay, 0.05, 0, background)
-        return background
+        return create_gradient_background(bg_config, width, height)
+        
     except Exception as e:
         logger.error(f"Minimalist background creation failed: {e}")
         return None
@@ -988,14 +538,18 @@ def create_corporate_background(colors, width, height):
             "direction": "diagonal"
         }
         background = create_gradient_background(bg_config, width, height)
+        
+        # Add subtle grid pattern
         grid_color = (80, 80, 80)
         grid_spacing = width // 20
         for x in range(0, width, grid_spacing):
             cv2.line(background, (x, 0), (x, height), grid_color, 1)
         for y in range(0, height, grid_spacing):
             cv2.line(background, (0, y), (width, y), grid_color, 1)
+        
         background = cv2.GaussianBlur(background, (3, 3), 1.0)
         return background
+        
     except Exception as e:
         logger.error(f"Corporate background creation failed: {e}")
         return None
@@ -1008,21 +562,8 @@ def create_nature_background(colors, width, height):
             "colors": ['#2d5016', '#3c6e1f', '#4caf50'],
             "direction": "vertical"
         }
-        background = create_gradient_background(bg_config, width, height)
-        import random
-        random.seed(42)
-        overlay = background.copy()
-        for _ in range(5):
-            center_x = random.randint(width//6, 5*width//6)
-            center_y = random.randint(height//6, 5*height//6)
-            axes_x = random.randint(width//20, width//6)
-            axes_y = random.randint(height//20, height//6)
-            angle = random.randint(0, 180)
-            color = (random.randint(40, 80), random.randint(120, 160), random.randint(30, 70))
-            cv2.ellipse(overlay, (center_x, center_y), (axes_x, axes_y), angle, 0, 360, color, -1)
-        cv2.addWeighted(background, 0.8, overlay, 0.2, 0, background)
-        background = cv2.GaussianBlur(background, (5, 5), 2.0)
-        return background
+        return create_gradient_background(bg_config, width, height)
+        
     except Exception as e:
         logger.error(f"Nature background creation failed: {e}")
         return None
@@ -1036,6 +577,8 @@ def create_artistic_background(colors, width, height):
             "direction": "diagonal"
         }
         background = create_gradient_background(bg_config, width, height)
+        
+        # Add artistic wave patterns
         import random
         random.seed(42)
         bgr_colors = []
@@ -1043,6 +586,7 @@ def create_artistic_background(colors, width, height):
             hex_color = color.lstrip('#')
             rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
             bgr_colors.append(rgb[::-1])
+        
         overlay = background.copy()
         for i in range(3):
             pts = []
@@ -1052,12 +596,35 @@ def create_artistic_background(colors, width, height):
             pts = np.array(pts, np.int32)
             color = bgr_colors[i % len(bgr_colors)]
             cv2.polylines(overlay, [pts], False, color, thickness=width//50)
+        
         cv2.addWeighted(background, 0.7, overlay, 0.3, 0, background)
         background = cv2.GaussianBlur(background, (3, 3), 1.0)
         return background
+        
     except Exception as e:
         logger.error(f"Artistic background creation failed: {e}")
         return None
 
-# ========== END OF FILE ==========
+def get_model_status():
+    """Get current model loading status"""
+    return "Models loaded in app.py - ready for processing"
+
+def validate_video_file(video_path):
+    """Validate video file format and basic properties"""
+    if not video_path or not os.path.exists(video_path):
+        return False, "Video file not found"
+    try:
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return False, "Cannot open video file"
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        if frame_count == 0:
+            return False, "Video appears to be empty"
+        cap.release()
+        return True, "Video file valid"
+    except Exception as e:
+        return False, f"Error validating video: {str(e)}"
 
+def get_available_backgrounds():
+    """Get list of available professional backgrounds"""
+    return {key: config["name"] for key, config in PROFESSIONAL_BACKGROUNDS.items()}