video_processor.py

"""
Core Video Processing Module - Enhanced with Temporal Consistency
VERSION: 2.0-temporal-enhanced
ROLLBACK: Set USE_TEMPORAL_ENHANCEMENT = False to revert to original behavior
"""

import os
import cv2
import numpy as np
import time
import logging
import threading
from typing import Optional, Tuple, Dict, Any, Callable, List
from pathlib import Path

# Import modular components
import app_config
import memory_manager
import progress_tracker
import exceptions

# Import utilities
from utilities import (
    segment_person_hq,
    refine_mask_hq,
    replace_background_hq,
    create_professional_background,
    PROFESSIONAL_BACKGROUNDS,
    validate_video_file
)

# ============================================================================
# VERSION CONTROL AND FEATURE FLAGS - EASY ROLLBACK
# ============================================================================

# ROLLBACK CONTROL: Set to False to use original functions
USE_TEMPORAL_ENHANCEMENT = True
USE_HAIR_DETECTION = True
USE_OPTICAL_FLOW_TRACKING = True
USE_ADAPTIVE_REFINEMENT = True

logger = logging.getLogger(__name__)

class CoreVideoProcessor:
    """
    ENHANCED: Core video processing pipeline with temporal consistency and fine-detail handling
    """
    
    def __init__(self, sam2_predictor: Any, matanyone_model: Any, 
                 config: app_config.ProcessingConfig, memory_mgr: memory_manager.MemoryManager):
        self.sam2_predictor = sam2_predictor
        self.matanyone_model = matanyone_model
        self.config = config
        self.memory_manager = memory_mgr
        
        # Processing state
        self.processing_active = False
        self.last_refined_mask = None
        self.frame_cache = {}
        
        # ENHANCED: Temporal consistency state
        self.mask_history = []  # Store recent masks for temporal smoothing
        self.optical_flow_data = None  # Previous frame for optical flow
        self.hair_regions_cache = {}  # Cache detected hair regions
        self.quality_scores_history = []  # Track quality over time
        
        # Statistics
        self.stats = {
            'videos_processed': 0,
            'total_frames_processed': 0,
            'total_processing_time': 0.0,
            'average_fps': 0.0,
            'failed_frames': 0,
            'successful_frames': 0,
            'cache_hits': 0,
            'segmentation_errors': 0,
            'refinement_errors': 0,
            'temporal_corrections': 0,  # NEW: Track temporal fixes
            'hair_detections': 0,       # NEW: Track hair detection success
            'flow_tracking_failures': 0  # NEW: Track optical flow issues
        }
        
        # Quality settings based on config
        self.quality_settings = config.get_quality_settings()
        
        logger.info("CoreVideoProcessor initialized")
        logger.info(f"Quality preset: {config.quality_preset}")
        logger.info(f"Quality settings: {self.quality_settings}")
        
        if USE_TEMPORAL_ENHANCEMENT:
            logger.info("ENHANCED: Temporal consistency enabled")
        if USE_HAIR_DETECTION:
            logger.info("ENHANCED: Hair detection enabled")
    
    def process_video(
        self,
        video_path: str,
        background_choice: str,
        custom_background_path: Optional[str] = None,
        progress_callback: Optional[Callable] = None,
        cancel_event: Optional[threading.Event] = None,
        preview_mask: bool = False,
        preview_greenscreen: bool = False
    ) -> Tuple[Optional[str], str]:
        """
        ENHANCED: Process video with temporal consistency and fine-detail handling
        """
        if self.processing_active:
            return None, "Processing already in progress"
        
        self.processing_active = True
        start_time = time.time()
        
        # ENHANCED: Reset temporal state for new video
        self._reset_temporal_state()
        
        try:
            # Validate input video
            is_valid, validation_msg = validate_video_file(video_path)
            if not is_valid:
                return None, f"Invalid video file: {validation_msg}"
            
            # Open video file
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                return None, "Could not open video file"
            
            # Get video properties
            video_info = self._get_video_info(cap)
            logger.info(f"Processing video: {video_info}")
            
            # Check memory requirements
            memory_check = self.memory_manager.can_process_video(
                video_info['width'], video_info['height']
            )
            
            if not memory_check['can_process']:
                cap.release()
                return None, f"Insufficient memory: {memory_check['recommendations']}"
            
            # Prepare background
            background = self.prepare_background(
                background_choice, custom_background_path, 
                video_info['width'], video_info['height']
            )
            
            if background is None:
                cap.release()
                return None, "Failed to prepare background"
            
            # Setup output video
            output_path = self._setup_output_video(video_info, preview_mask, preview_greenscreen)
            out = self._create_video_writer(output_path, video_info)
            
            if out is None:
                cap.release()
                return None, "Could not create output video writer"
            
            # ENHANCED: Process video frames with temporal consistency
            result = self._process_video_frames_enhanced(
                cap, out, background, video_info, 
                progress_callback, cancel_event, 
                preview_mask, preview_greenscreen
            )
            
            # Cleanup
            cap.release()
            out.release()
            
            if result['success']:
                # Update statistics
                processing_time = time.time() - start_time
                self._update_processing_stats(video_info, processing_time, result)
                
                success_msg = (
                    f"Processing completed successfully!\n"
                    f"Processed: {result['successful_frames']}/{result['total_frames']} frames\n"
                    f"Time: {processing_time:.1f}s\n"
                    f"Average FPS: {result['total_frames'] / processing_time:.1f}\n"
                    f"Temporal corrections: {self.stats['temporal_corrections']}\n"
                    f"Hair detections: {self.stats['hair_detections']}\n"
                    f"Background: {background_choice}"
                )
                
                return output_path, success_msg
            else:
                # Clean up failed output
                try:
                    os.remove(output_path)
                except:
                    pass
                return None, result['error_message']
                
        except Exception as e:
            logger.error(f"Video processing failed: {e}")
            return None, f"Processing failed: {str(e)}"
        
        finally:
            self.processing_active = False
    
    def _reset_temporal_state(self):
        """ENHANCED: Reset temporal consistency state"""
        self.mask_history.clear()
        self.optical_flow_data = None
        self.hair_regions_cache.clear()
        self.quality_scores_history.clear()
        self.last_refined_mask = None
        self.stats['temporal_corrections'] = 0
        self.stats['hair_detections'] = 0
        self.stats['flow_tracking_failures'] = 0
    
    def _get_video_info(self, cap: cv2.VideoCapture) -> Dict[str, Any]:
        """Extract comprehensive video information"""
        return {
            'fps': cap.get(cv2.CAP_PROP_FPS),
            'total_frames': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
            'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
            'duration': cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS),
            'codec': int(cap.get(cv2.CAP_PROP_FOURCC))
        }
    
    def _setup_output_video(self, video_info: Dict[str, Any], 
                           preview_mask: bool, preview_greenscreen: bool) -> str:
        """Setup output video path"""
        timestamp = int(time.time())
        
        if preview_mask:
            filename = f"mask_preview_{timestamp}.mp4"
        elif preview_greenscreen:
            filename = f"greenscreen_preview_{timestamp}.mp4"
        else:
            filename = f"processed_video_{timestamp}.mp4"
        
        return os.path.join(self.config.temp_dir, filename)
    
    def _create_video_writer(self, output_path: str, 
                           video_info: Dict[str, Any]) -> Optional[cv2.VideoWriter]:
        """Create video writer with optimal settings"""
        try:
            # Choose codec based on quality settings
            if self.config.output_quality == 'high':
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            else:
                fourcc = cv2.VideoWriter_fourcc(*'XVID')
            
            writer = cv2.VideoWriter(
                output_path, 
                fourcc, 
                video_info['fps'],
                (video_info['width'], video_info['height'])
            )
            
            if not writer.isOpened():
                logger.error("Failed to open video writer")
                return None
                
            return writer
            
        except Exception as e:
            logger.error(f"Error creating video writer: {e}")
            return None
    
    def _process_video_frames_enhanced(
        self,
        cap: cv2.VideoCapture,
        out: cv2.VideoWriter,
        background: np.ndarray,
        video_info: Dict[str, Any],
        progress_callback: Optional[Callable],
        cancel_event: Optional[threading.Event],
        preview_mask: bool,
        preview_greenscreen: bool
    ) -> Dict[str, Any]:
        """ENHANCED: Process all video frames with temporal consistency"""
        
        # Initialize progress tracking
        prog_tracker = progress_tracker.ProgressTracker(
            total_frames=video_info['total_frames'],
            callback=progress_callback,
            track_performance=True
        )
        
        frame_count = 0
        successful_frames = 0
        failed_frames = 0
        
        # Reset enhanced state
        self._reset_temporal_state()
        
        try:
            prog_tracker.set_stage("Processing frames with temporal enhancement")
            
            while True:
                # Check for cancellation
                if cancel_event and cancel_event.is_set():
                    return {
                        'success': False,
                        'error_message': 'Processing cancelled by user',
                        'total_frames': frame_count,
                        'successful_frames': successful_frames,
                        'failed_frames': failed_frames
                    }
                
                # Read frame
                ret, frame = cap.read()
                if not ret:
                    break
                
                try:
                    # Update progress
                    prog_tracker.update(frame_count, "Processing frame with temporal consistency")
                    
                    # ENHANCED: Process frame with temporal consistency
                    if USE_TEMPORAL_ENHANCEMENT:
                        processed_frame = self._process_single_frame_enhanced(
                            frame, background, frame_count, 
                            preview_mask, preview_greenscreen
                        )
                    else:
                        processed_frame = self._process_single_frame_original(
                            frame, background, frame_count, 
                            preview_mask, preview_greenscreen
                        )
                    
                    # Write processed frame
                    out.write(processed_frame)
                    successful_frames += 1
                    
                    # Memory management
                    if frame_count % self.config.memory_cleanup_interval == 0:
                        self.memory_manager.auto_cleanup_if_needed()
                    
                except Exception as frame_error:
                    logger.warning(f"Frame {frame_count} processing failed: {frame_error}")
                    
                    # Write original frame as fallback
                    out.write(frame)
                    failed_frames += 1
                    self.stats['failed_frames'] += 1
                
                frame_count += 1
                
                # Skip frames if configured (for performance)
                if self.config.frame_skip > 1:
                    for _ in range(self.config.frame_skip - 1):
                        ret, _ = cap.read()
                        if not ret:
                            break
                        frame_count += 1
            
            # Finalize progress tracking
            final_stats = prog_tracker.finalize()
            
            return {
                'success': successful_frames > 0,
                'error_message': f'No frames processed successfully' if successful_frames == 0 else '',
                'total_frames': frame_count,
                'successful_frames': successful_frames,
                'failed_frames': failed_frames,
                'processing_stats': final_stats
            }
            
        except Exception as e:
            logger.error(f"Frame processing loop failed: {e}")
            return {
                'success': False,
                'error_message': f'Frame processing failed: {str(e)}',
                'total_frames': frame_count,
                'successful_frames': successful_frames,
                'failed_frames': failed_frames
            }
    
    def _process_single_frame_enhanced(
        self,
        frame: np.ndarray,
        background: np.ndarray,
        frame_number: int,
        preview_mask: bool,
        preview_greenscreen: bool
    ) -> np.ndarray:
        """ENHANCED: Process a single video frame with temporal consistency"""
        
        try:
            # Person segmentation
            mask = self._segment_person_enhanced(frame, frame_number)
            
            # ENHANCED: Detect hair and fine details
            if USE_HAIR_DETECTION:
                hair_regions = self._detect_hair_regions(frame, mask, frame_number)
            else:
                hair_regions = None
            
            # ENHANCED: Apply temporal consistency
            if USE_TEMPORAL_ENHANCEMENT and len(self.mask_history) > 0:
                mask = self._apply_temporal_consistency_enhanced(frame, mask, frame_number)
            
            # ENHANCED: Adaptive mask refinement based on frame content
            if USE_ADAPTIVE_REFINEMENT:
                refined_mask = self._adaptive_mask_refinement(frame, mask, frame_number, hair_regions)
            else:
                refined_mask = self._refine_mask_original(frame, mask, frame_number)
            
            # Store mask in history for temporal consistency
            self._update_mask_history(refined_mask)
            
            # Generate output based on mode
            if preview_mask:
                return self._create_mask_preview_enhanced(frame, refined_mask, hair_regions)
            elif preview_greenscreen:
                return self._create_greenscreen_preview(frame, refined_mask)
            else:
                return self._replace_background_enhanced(frame, refined_mask, background, hair_regions)
                
        except Exception as e:
            logger.warning(f"Enhanced single frame processing failed: {e}")
            # Fallback to original processing
            return self._process_single_frame_original(frame, background, frame_number, preview_mask, preview_greenscreen)
    
    def _detect_hair_regions(self, frame: np.ndarray, mask: np.ndarray, frame_number: int) -> Optional[np.ndarray]:
        """ENHANCED: Detect hair and fine detail regions automatically"""
        try:
            # Check cache first
            if frame_number in self.hair_regions_cache:
                self.stats['cache_hits'] += 1
                return self.hair_regions_cache[frame_number]
            
            # Convert frame to different color spaces for better hair detection
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # Method 1: Texture-based hair detection
            # Hair typically has high frequency texture
            laplacian = cv2.Laplacian(gray, cv2.CV_64F)
            texture_strength = np.abs(laplacian)
            
            # Method 2: Color-based hair detection
            # Hair is typically in darker hue ranges
            hair_hue_mask = ((hsv[:,:,0] >= 0) & (hsv[:,:,0] <= 30)) | \
                           ((hsv[:,:,0] >= 150) & (hsv[:,:,0] <= 180))
            hair_value_mask = hsv[:,:,2] < 100  # Darker regions
            
            # Combine texture and color information
            hair_probability = np.zeros_like(gray, dtype=np.float32)
            
            # High texture regions
            texture_norm = (texture_strength - texture_strength.min()) / (texture_strength.max() - texture_strength.min() + 1e-8)
            hair_probability += texture_norm * 0.6
            
            # Color-based probability
            color_prob = (hair_hue_mask.astype(np.float32) * hair_value_mask.astype(np.float32))
            hair_probability += color_prob * 0.4
            
            # Only consider regions near the mask boundary (where hair typically is)
            mask_boundary = self._get_mask_boundary_region(mask, boundary_width=20)
            hair_probability *= mask_boundary
            
            # Threshold to get hair regions
            hair_threshold = np.percentile(hair_probability[hair_probability > 0], 75)
            hair_regions = (hair_probability > hair_threshold).astype(np.uint8)
            
            # Clean up hair regions
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
            hair_regions = cv2.morphologyEx(hair_regions, cv2.MORPH_CLOSE, kernel)
            
            # Cache the result
            self.hair_regions_cache[frame_number] = hair_regions
            
            # Update stats if hair was detected
            if np.any(hair_regions):
                self.stats['hair_detections'] += 1
                logger.debug(f"Hair regions detected in frame {frame_number}")
            
            return hair_regions
            
        except Exception as e:
            logger.warning(f"Hair detection failed for frame {frame_number}: {e}")
            return None
    
    def _get_mask_boundary_region(self, mask: np.ndarray, boundary_width: int = 20) -> np.ndarray:
        """Get region around mask boundary where hair/fine details are likely"""
        try:
            # Create dilated and eroded versions of mask
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (boundary_width, boundary_width))
            dilated = cv2.dilate(mask, kernel, iterations=1)
            eroded = cv2.erode(mask, kernel, iterations=1)
            
            # Boundary region is dilated minus eroded
            boundary_region = ((dilated > 0) & (eroded == 0)).astype(np.float32)
            
            return boundary_region
            
        except Exception as e:
            logger.warning(f"Boundary region detection failed: {e}")
            return np.ones_like(mask, dtype=np.float32)
    
    def _apply_temporal_consistency_enhanced(self, frame: np.ndarray, current_mask: np.ndarray, frame_number: int) -> np.ndarray:
        """ENHANCED: Apply temporal consistency using optical flow and history"""
        try:
            if len(self.mask_history) == 0:
                return current_mask
            
            previous_mask = self.mask_history[-1]
            
            # Method 1: Optical flow-based consistency
            if USE_OPTICAL_FLOW_TRACKING and self.optical_flow_data is not None:
                try:
                    flow_corrected_mask = self._apply_optical_flow_consistency(
                        frame, current_mask, previous_mask
                    )
                    
                    # Blend flow-corrected with current mask
                    alpha = 0.7  # Weight for current mask
                    beta = 0.3   # Weight for flow-corrected mask
                    
                    blended_mask = cv2.addWeighted(
                        current_mask.astype(np.float32), alpha,
                        flow_corrected_mask.astype(np.float32), beta, 0
                    ).astype(np.uint8)
                    
                    self.stats['temporal_corrections'] += 1
                    
                except Exception as e:
                    logger.debug(f"Optical flow consistency failed: {e}")
                    self.stats['flow_tracking_failures'] += 1
                    blended_mask = current_mask
            else:
                blended_mask = current_mask
            
            # Method 2: Multi-frame temporal smoothing
            if len(self.mask_history) >= 3:
                # Use weighted average of recent masks
                weights = [0.5, 0.3, 0.2]  # Current, previous, before previous
                masks_to_blend = [blended_mask] + self.mask_history[-2:]
                
                temporal_mask = np.zeros_like(blended_mask, dtype=np.float32)
                for mask, weight in zip(masks_to_blend, weights):
                    temporal_mask += mask.astype(np.float32) * weight
                
                blended_mask = np.clip(temporal_mask, 0, 255).astype(np.uint8)
            
            # Method 3: Edge-aware temporal filtering
            blended_mask = self._temporal_edge_filtering(frame, blended_mask, current_mask)
            
            return blended_mask
            
        except Exception as e:
            logger.warning(f"Temporal consistency failed: {e}")
            return current_mask
    
    def _apply_optical_flow_consistency(self, current_frame: np.ndarray, 
                                      current_mask: np.ndarray, previous_mask: np.ndarray) -> np.ndarray:
        """Apply optical flow to warp previous mask to current frame"""
        try:
            # Convert frames to grayscale for optical flow
            current_gray = cv2.cvtColor(current_frame, cv2.COLOR_BGR2GRAY)
            previous_gray = self.optical_flow_data
            
            # Calculate dense optical flow
            flow = cv2.calcOpticalFlowPyrLK(previous_gray, current_gray, None, None)
            
            # Warp previous mask using optical flow
            h, w = previous_mask.shape
            flow_map = np.zeros((h, w, 2), dtype=np.float32)
            
            # Create flow field
            y_coords, x_coords = np.mgrid[0:h, 0:w]
            flow_map[:, :, 0] = x_coords + flow[0] if flow[0] is not None else x_coords
            flow_map[:, :, 1] = y_coords + flow[1] if flow[1] is not None else y_coords
            
            # Warp previous mask
            warped_mask = cv2.remap(previous_mask, flow_map, None, cv2.INTER_LINEAR)
            
            return warped_mask
            
        except Exception as e:
            logger.debug(f"Optical flow warping failed: {e}")
            return previous_mask
    
    def _temporal_edge_filtering(self, frame: np.ndarray, blended_mask: np.ndarray, current_mask: np.ndarray) -> np.ndarray:
        """Apply edge-aware temporal filtering"""
        try:
            # Detect edges in current frame
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            
            # In edge regions, favor the current mask (more responsive)
            # In smooth regions, favor the blended mask (more stable)
            edge_weight = cv2.GaussianBlur(edges.astype(np.float32), (5, 5), 1.0) / 255.0
            
            filtered_mask = (current_mask.astype(np.float32) * edge_weight + 
                           blended_mask.astype(np.float32) * (1 - edge_weight))
            
            return np.clip(filtered_mask, 0, 255).astype(np.uint8)
            
        except Exception as e:
            logger.warning(f"Temporal edge filtering failed: {e}")
            return blended_mask
    
    def _adaptive_mask_refinement(self, frame: np.ndarray, mask: np.ndarray, 
                                frame_number: int, hair_regions: Optional[np.ndarray]) -> np.ndarray:
        """ENHANCED: Adaptive mask refinement based on content analysis"""
        try:
            # Determine refinement strategy based on frame content
            refinement_needed = self._assess_refinement_needs(frame, mask, hair_regions)
            
            if refinement_needed['hair_refinement'] and hair_regions is not None:
                # Special handling for hair regions
                mask = self._refine_hair_regions(frame, mask, hair_regions)
            
            if refinement_needed['edge_refinement']:
                # Enhanced edge refinement
                mask = self._enhanced_edge_refinement(frame, mask)
            
            if refinement_needed['temporal_refinement']:
                # Apply temporal-aware refinement
                mask = self._temporal_aware_refinement(frame, mask, frame_number)
            
            # Standard refinement if needed
            if self._should_refine_mask(frame_number):
                if self.matanyone_model is not None and self.quality_settings.get('edge_refinement', True):
                    mask = refine_mask_hq(frame, mask, self.matanyone_model)
                else:
                    mask = self._fallback_mask_refinement_enhanced(mask)
            
            return mask
            
        except Exception as e:
            logger.warning(f"Adaptive mask refinement failed: {e}")
            return self._refine_mask_original(frame, mask, frame_number)
    
    def _assess_refinement_needs(self, frame: np.ndarray, mask: np.ndarray, 
                               hair_regions: Optional[np.ndarray]) -> Dict[str, bool]:
        """Assess what type of refinement is needed for this frame"""
        try:
            needs = {
                'hair_refinement': False,
                'edge_refinement': False,
                'temporal_refinement': False
            }
            
            # Check if hair refinement is needed
            if hair_regions is not None and np.any(hair_regions):
                needs['hair_refinement'] = True
            
            # Check edge quality
            edges = cv2.Canny(mask, 50, 150)
            edge_density = np.sum(edges > 0) / (mask.shape[0] * mask.shape[1])
            if edge_density > 0.1:  # High edge density suggests rough boundaries
                needs['edge_refinement'] = True
            
            # Check temporal consistency needs
            if len(self.mask_history) > 0:
                prev_mask = self.mask_history[-1]
                diff = cv2.absdiff(mask, prev_mask)
                change_ratio = np.sum(diff > 50) / (mask.shape[0] * mask.shape[1])
                if change_ratio > 0.15:  # High change suggests temporal inconsistency
                    needs['temporal_refinement'] = True
            
            return needs
            
        except Exception as e:
            logger.warning(f"Refinement assessment failed: {e}")
            return {'hair_refinement': False, 'edge_refinement': True, 'temporal_refinement': False}
    
    def _refine_hair_regions(self, frame: np.ndarray, mask: np.ndarray, hair_regions: np.ndarray) -> np.ndarray:
        """Special refinement for hair and fine detail regions"""
        try:
            # Create a more aggressive mask for hair regions
            hair_mask = hair_regions > 0
            
            # Use different thresholding for hair areas
            refined_mask = mask.copy()
            
            # In hair regions, use lower threshold (include more pixels)
            hair_area_values = mask[hair_mask]
            if len(hair_area_values) > 0:
                hair_threshold = max(100, np.percentile(hair_area_values, 25))  # Lower threshold for hair
                refined_mask[hair_mask] = np.where(mask[hair_mask] > hair_threshold, 255, 0)
            
            # Apply morphological closing to connect hair strands
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
            refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_CLOSE, kernel)
            
            return refined_mask
            
        except Exception as e:
            logger.warning(f"Hair region refinement failed: {e}")
            return mask
    
    def _enhanced_edge_refinement(self, frame: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Enhanced edge refinement using image gradients"""
        try:
            # Use bilateral filter to preserve edges while smoothing
            refined = cv2.bilateralFilter(mask, 9, 75, 75)
            
            # Edge-guided smoothing
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            
            # In edge areas, preserve original mask more
            edge_weight = cv2.GaussianBlur(edges.astype(np.float32), (3, 3), 1.0) / 255.0
            edge_weight = np.clip(edge_weight * 2, 0, 1)  # Amplify edge influence
            
            final_mask = (mask.astype(np.float32) * edge_weight + 
                         refined.astype(np.float32) * (1 - edge_weight))
            
            return np.clip(final_mask, 0, 255).astype(np.uint8)
            
        except Exception as e:
            logger.warning(f"Enhanced edge refinement failed: {e}")
            return mask
    
    def _temporal_aware_refinement(self, frame: np.ndarray, mask: np.ndarray, frame_number: int) -> np.ndarray:
        """Temporal-aware refinement considering motion and stability"""
        try:
            if len(self.mask_history) == 0:
                return mask
            
            # Calculate motion between frames
            if self.optical_flow_data is not None:
                current_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                motion_magnitude = cv2.absdiff(current_gray, self.optical_flow_data)
                motion_mask = motion_magnitude > 10  # Areas with motion
                
                # In high-motion areas, trust current mask more
                # In low-motion areas, use temporal smoothing
                prev_mask = self.mask_history[-1]
                
                motion_weight = cv2.GaussianBlur(motion_mask.astype(np.float32), (5, 5), 1.0)
                motion_weight = np.clip(motion_weight, 0.3, 1.0)  # Don't completely ignore temporal info
                
                temporal_mask = (mask.astype(np.float32) * motion_weight + 
                               prev_mask.astype(np.float32) * (1 - motion_weight))
                
                return np.clip(temporal_mask, 0, 255).astype(np.uint8)
            
            return mask
            
        except Exception as e:
            logger.warning(f"Temporal-aware refinement failed: {e}")
            return mask
    
    def _update_mask_history(self, mask: np.ndarray):
        """Update mask history for temporal consistency"""
        self.mask_history.append(mask.copy())
        
        # Keep only recent history (limit memory usage)
        max_history = 5
        if len(self.mask_history) > max_history:
            self.mask_history.pop(0)
    
    def _create_mask_preview_enhanced(self, frame: np.ndarray, mask: np.ndarray, 
                                    hair_regions: Optional[np.ndarray]) -> np.ndarray:
        """ENHANCED: Create mask visualization with hair regions highlighted"""
        try:
            # Create colored mask overlay
            mask_colored = np.zeros_like(frame)
            mask_colored[:, :, 1] = mask  # Green channel for person
            
            # Highlight hair regions in blue if available
            if hair_regions is not None:
                mask_colored[:, :, 2] = np.maximum(mask_colored[:, :, 2], hair_regions * 255)
            
            # Blend with original frame
            alpha = 0.6
            preview = cv2.addWeighted(frame, 1-alpha, mask_colored, alpha, 0)
            
            return preview
            
        except Exception as e:
            logger.warning(f"Enhanced mask preview creation failed: {e}")
            return self._create_mask_preview_original(frame, mask)
    
    def _replace_background_enhanced(self, frame: np.ndarray, mask: np.ndarray, 
                                   background: np.ndarray, hair_regions: Optional[np.ndarray]) -> np.ndarray:
        """ENHANCED: Replace background with special handling for hair regions"""
        try:
            # Standard background replacement
            result = replace_background_hq(frame, mask, background)
            
            # If hair regions detected, apply additional processing
            if hair_regions is not None and np.any(hair_regions):
                result = self._enhance_hair_compositing(frame, mask, background, hair_regions, result)
            
            return result
            
        except Exception as e:
            logger.warning(f"Enhanced background replacement failed: {e}")
            return replace_background_hq(frame, mask, background)
    
    def _enhance_hair_compositing(self, frame: np.ndarray, mask: np.ndarray, 
                                background: np.ndarray, hair_regions: np.ndarray, 
                                initial_result: np.ndarray) -> np.ndarray:
        """Enhanced compositing specifically for hair regions"""
        try:
            # In hair regions, use softer alpha blending
            hair_mask = hair_regions > 0
            
            if np.any(hair_mask):
                # Create soft alpha for hair regions
                hair_alpha = cv2.GaussianBlur((hair_regions * mask / 255.0).astype(np.float32), (3, 3), 1.0)
                hair_alpha = np.clip(hair_alpha, 0, 1)
                
                # Apply softer blending only in hair regions
                for c in range(3):
                    channel_blend = (frame[:, :, c].astype(np.float32) * hair_alpha + 
                                   background[:, :, c].astype(np.float32) * (1 - hair_alpha))
                    
                    initial_result[:, :, c] = np.where(
                        hair_mask,
                        np.clip(channel_blend, 0, 255).astype(np.uint8),
                        initial_result[:, :, c]
                    )
            
            return initial_result
            
        except Exception as e:
            logger.warning(f"Hair compositing enhancement failed: {e}")
            return initial_result
    
    # ============================================================================
    # ORIGINAL FUNCTIONS PRESERVED FOR ROLLBACK
    # ============================================================================
    
    def _process_single_frame_original(
        self,
        frame: np.ndarray,
        background: np.ndarray,
        frame_number: int,
        preview_mask: bool,
        preview_greenscreen: bool
    ) -> np.ndarray:
        """ORIGINAL: Process a single video frame (preserved for rollback)"""
        
        try:
            # Person segmentation
            mask = self._segment_person(frame, frame_number)
            
            # Mask refinement (keyframe-based for performance)
            if self._should_refine_mask(frame_number):
                refined_mask = self._refine_mask_original(frame, mask, frame_number)
                self.last_refined_mask = refined_mask.copy()
            else:
                # Use temporal consistency with previous refined mask
                refined_mask = self._apply_temporal_consistency_original(mask, frame_number)
            
            # Generate output based on mode
            if preview_mask:
                return self._create_mask_preview_original(frame, refined_mask)
            elif preview_greenscreen:
                return self._create_greenscreen_preview(frame, refined_mask)
            else:
                return self._replace_background(frame, refined_mask, background)
                
        except Exception as e:
            logger.warning(f"Single frame processing failed: {e}")
            raise
    
    def _segment_person(self, frame: np.ndarray, frame_number: int) -> np.ndarray:
        """Perform person segmentation"""
        try:
            mask = segment_person_hq(frame, self.sam2_predictor)
            
            if mask is None or mask.size == 0:
                raise exceptions.SegmentationError(frame_number, "Segmentation returned empty mask")
            
            # Store current frame for optical flow (if enhanced mode enabled)
            if USE_OPTICAL_FLOW_TRACKING:
                current_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                self.optical_flow_data = current_gray
            
            return mask
            
        except Exception as e:
            self.stats['segmentation_errors'] += 1
            raise exceptions.SegmentationError(frame_number, f"Segmentation failed: {str(e)}")
    
    def _segment_person_enhanced(self, frame: np.ndarray, frame_number: int) -> np.ndarray:
        """ENHANCED: Perform person segmentation with improvements"""
        try:
            mask = segment_person_hq(frame, self.sam2_predictor)
            
            if mask is None or mask.size == 0:
                raise exceptions.SegmentationError(frame_number, "Segmentation returned empty mask")
            
            # Store current frame for optical flow
            current_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            self.optical_flow_data = current_gray
            
            return mask
            
        except Exception as e:
            self.stats['segmentation_errors'] += 1
            raise exceptions.SegmentationError(frame_number, f"Enhanced segmentation failed: {str(e)}")
    
    def _should_refine_mask(self, frame_number: int) -> bool:
        """Determine if mask should be refined for this frame"""
        # Refine on keyframes or if no previous refined mask exists
        return (
            frame_number % self.quality_settings['keyframe_interval'] == 0 or 
            self.last_refined_mask is None or
            not self.quality_settings.get('temporal_consistency', True)
        )
    
    def _refine_mask_original(self, frame: np.ndarray, mask: np.ndarray, frame_number: int) -> np.ndarray:
        """ORIGINAL: Refine mask using MatAnyone or fallback methods"""
        try:
            if self.matanyone_model is not None and self.quality_settings.get('edge_refinement', True):
                refined_mask = refine_mask_hq(frame, mask, self.matanyone_model)
            else:
                # Fallback refinement using OpenCV operations
                refined_mask = self._fallback_mask_refinement(mask)
            
            return refined_mask
            
        except Exception as e:
            self.stats['refinement_errors'] += 1
            logger.warning(f"Mask refinement failed for frame {frame_number}: {e}")
            # Return original mask as fallback
            return mask
    
    def _fallback_mask_refinement(self, mask: np.ndarray) -> np.ndarray:
        """ORIGINAL: Fallback mask refinement using basic OpenCV operations"""
        try:
            # Morphological operations to clean up mask
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
            refined = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
            refined = cv2.morphologyEx(refined, cv2.MORPH_OPEN, kernel)
            
            # Smooth edges
            refined = cv2.GaussianBlur(refined, (3, 3), 1.0)
            
            return refined
            
        except Exception as e:
            logger.warning(f"Fallback mask refinement failed: {e}")
            return mask
    
    def _fallback_mask_refinement_enhanced(self, mask: np.ndarray) -> np.ndarray:
        """ENHANCED: Improved fallback mask refinement"""
        try:
            # More aggressive morphological operations
            kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
            kernel_large = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
            
            # Remove small noise
            refined = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_small)
            # Fill gaps
            refined = cv2.morphologyEx(refined, cv2.MORPH_CLOSE, kernel_large)
            
            # Edge smoothing with bilateral filter instead of Gaussian
            refined = cv2.bilateralFilter(refined, 9, 75, 75)
            
            return refined
            
        except Exception as e:
            logger.warning(f"Enhanced fallback mask refinement failed: {e}")
            return mask
    
    def _apply_temporal_consistency_original(self, current_mask: np.ndarray, frame_number: int) -> np.ndarray:
        """ORIGINAL: Apply temporal consistency using previous refined mask"""
        if self.last_refined_mask is None or not self.quality_settings.get('temporal_consistency', True):
            return current_mask
        
        try:
            # Blend current mask with previous refined mask
            alpha = 0.7  # Weight for current mask
            beta = 0.3   # Weight for previous mask
            
            # Ensure masks have same shape
            if current_mask.shape != self.last_refined_mask.shape:
                last_mask = cv2.resize(self.last_refined_mask, 
                                     (current_mask.shape[1], current_mask.shape[0]))
            else:
                last_mask = self.last_refined_mask
            
            # Weighted blend
            blended_mask = cv2.addWeighted(current_mask, alpha, last_mask, beta, 0)
            
            # Apply slight smoothing for temporal stability
            blended_mask = cv2.GaussianBlur(blended_mask, (3, 3), 0.5)
            
            return blended_mask
            
        except Exception as e:
            logger.warning(f"Temporal consistency application failed: {e}")
            return current_mask
    
    def _create_mask_preview_original(self, frame: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """ORIGINAL: Create mask visualization preview"""
        try:
            # Create colored mask overlay
            mask_colored = np.zeros_like(frame)
            mask_colored[:, :, 1] = mask  # Green channel for person
            
            # Blend with original frame
            alpha = 0.6
            preview = cv2.addWeighted(frame, 1-alpha, mask_colored, alpha, 0)
            
            return preview
            
        except Exception as e:
            logger.warning(f"Mask preview creation failed: {e}")
            return frame
    
    def _create_greenscreen_preview(self, frame: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Create green screen preview"""
        try:
            # Create pure green background
            green_bg = np.zeros_like(frame)
            green_bg[:, :] = [0, 255, 0]  # Pure green in BGR
            
            # Apply mask
            mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) if len(mask.shape) == 2 else mask
            mask_norm = mask_3ch.astype(np.float32) / 255.0
            
            result = (frame * mask_norm + green_bg * (1 - mask_norm)).astype(np.uint8)
            
            return result
            
        except Exception as e:
            logger.warning(f"Greenscreen preview creation failed: {e}")
            return frame
    
    def _replace_background(self, frame: np.ndarray, mask: np.ndarray, background: np.ndarray) -> np.ndarray:
        """Replace background using the refined mask"""
        try:
            result = replace_background_hq(frame, mask, background)
            return result
            
        except Exception as e:
            logger.warning(f"Background replacement failed: {e}")
            return frame
    
    def prepare_background(
        self,
        background_choice: str,
        custom_background_path: Optional[str],
        width: int,
        height: int
    ) -> Optional[np.ndarray]:
        """Prepare background image for processing (unchanged)"""
        try:
            if background_choice == "custom" and custom_background_path:
                if not os.path.exists(custom_background_path):
                    raise exceptions.BackgroundProcessingError("custom", f"File not found: {custom_background_path}")
                
                background = cv2.imread(custom_background_path)
                if background is None:
                    raise exceptions.BackgroundProcessingError("custom", "Could not read custom background image")
                
                logger.info(f"Loaded custom background: {custom_background_path}")
                
            else:
                # Use professional background
                if background_choice not in PROFESSIONAL_BACKGROUNDS:
                    raise exceptions.BackgroundProcessingError(background_choice, "Unknown professional background")
                
                bg_config = PROFESSIONAL_BACKGROUNDS[background_choice]
                background = create_professional_background(bg_config, width, height)
                
                logger.info(f"Generated professional background: {background_choice}")
            
            # Resize to match video dimensions
            if background.shape[:2] != (height, width):
                background = cv2.resize(background, (width, height), interpolation=cv2.INTER_LANCZOS4)
            
            # Validate background
            if background is None or background.size == 0:
                raise exceptions.BackgroundProcessingError(background_choice, "Background image is empty")
            
            return background
            
        except Exception as e:
            if isinstance(e, exceptions.BackgroundProcessingError):
                logger.error(str(e))
                return None
            else:
                logger.error(f"Unexpected error preparing background: {e}")
                return None
    
    def _update_processing_stats(self, video_info: Dict[str, Any], 
                               processing_time: float, result: Dict[str, Any]):
        """Update processing statistics"""
        self.stats['videos_processed'] += 1
        self.stats['total_frames_processed'] += result['successful_frames']
        self.stats['total_processing_time'] += processing_time
        self.stats['successful_frames'] += result['successful_frames']
        self.stats['failed_frames'] += result['failed_frames']
        
        # Calculate average FPS across all processing
        if self.stats['total_processing_time'] > 0:
            self.stats['average_fps'] = self.stats['total_frames_processed'] / self.stats['total_processing_time']
    
    def get_processing_capabilities(self) -> Dict[str, Any]:
        """Get current processing capabilities"""
        capabilities = {
            'sam2_available': self.sam2_predictor is not None,
            'matanyone_available': self.matanyone_model is not None,
            'quality_preset': self.config.quality_preset,
            'supports_temporal_consistency': self.quality_settings.get('temporal_consistency', False),
            'supports_edge_refinement': self.quality_settings.get('edge_refinement', False),
            'keyframe_interval': self.quality_settings['keyframe_interval'],
            'max_resolution': self.config.get_resolution_limits(),
            'supported_formats': ['.mp4', '.avi', '.mov', '.mkv'],
            'memory_limit_gb': self.memory_manager.memory_limit_gb
        }
        
        # Add enhanced capabilities
        if USE_TEMPORAL_ENHANCEMENT:
            capabilities.update({
                'temporal_enhancement': True,
                'hair_detection': USE_HAIR_DETECTION,
                'optical_flow_tracking': USE_OPTICAL_FLOW_TRACKING,
                'adaptive_refinement': USE_ADAPTIVE_REFINEMENT
            })
        
        return capabilities
    
    def get_status(self) -> Dict[str, Any]:
        """Get current processor status"""
        status = {
            'processing_active': self.processing_active,
            'models_available': {
                'sam2': self.sam2_predictor is not None,
                'matanyone': self.matanyone_model is not None
            },
            'quality_settings': self.quality_settings,
            'statistics': self.stats.copy(),
            'cache_size': len(self.frame_cache),
            'memory_usage': self.memory_manager.get_memory_usage(),
            'capabilities': self.get_processing_capabilities()
        }
        
        # Add enhanced status
        if USE_TEMPORAL_ENHANCEMENT:
            status.update({
                'mask_history_length': len(self.mask_history),
                'hair_cache_size': len(self.hair_regions_cache),
                'optical_flow_active': self.optical_flow_data is not None
            })
        
        return status
    
    def optimize_for_video(self, video_info: Dict[str, Any]) -> Dict[str, Any]:
        """Optimize settings for specific video characteristics"""
        optimizations = {
            'original_settings': self.quality_settings.copy(),
            'optimizations_applied': []
        }
        
        try:
            # High resolution video optimizations
            if video_info['width'] * video_info['height'] > 1920 * 1080:
                if self.quality_settings['keyframe_interval'] < 10:
                    self.quality_settings['keyframe_interval'] = 10
                    optimizations['optimizations_applied'].append('increased_keyframe_interval_for_high_res')
            
            # Long video optimizations
            if video_info['duration'] > 300:  # 5 minutes
                if self.config.memory_cleanup_interval > 20:
                    self.config.memory_cleanup_interval = 20
                    optimizations['optimizations_applied'].append('increased_memory_cleanup_frequency')
            
            # Low FPS video optimizations
            if video_info['fps'] < 15:
                self.quality_settings['temporal_consistency'] = False
                optimizations['optimizations_applied'].append('disabled_temporal_consistency_for_low_fps')
            
            # Memory-constrained optimizations
            memory_usage = self.memory_manager.get_memory_usage()
            memory_pressure = self.memory_manager.check_memory_pressure()
            
            if memory_pressure['under_pressure']:
                self.quality_settings['edge_refinement'] = False
                self.quality_settings['keyframe_interval'] = max(self.quality_settings['keyframe_interval'], 15)
                optimizations['optimizations_applied'].extend([
                    'disabled_edge_refinement_for_memory',
                    'increased_keyframe_interval_for_memory'
                ])
            
            optimizations['final_settings'] = self.quality_settings.copy()
            
            if optimizations['optimizations_applied']:
                logger.info(f"Applied video optimizations: {optimizations['optimizations_applied']}")
            
            return optimizations
            
        except Exception as e:
            logger.warning(f"Video optimization failed: {e}")
            return optimizations
    
    def reset_cache(self):
        """Reset frame cache and temporal state"""
        self.frame_cache.clear()
        self.last_refined_mask = None
        self.stats['cache_hits'] = 0
        self._reset_temporal_state()
        logger.debug("Frame cache and temporal state reset")
    
    def cleanup(self):
        """Clean up processor resources"""
        try:
            self.reset_cache()
            self.processing_active = False
            logger.info("CoreVideoProcessor cleanup completed")
        except Exception as e:
            logger.warning(f"Error during cleanup: {e}")