video_processor.py

"""
Core Video Processing Module
Handles the main video processing pipeline, frame processing, and background replacement
"""

import os
import cv2
import numpy as np
import time
import logging
import threading
from typing import Optional, Tuple, Dict, Any, Callable
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
)

logger = logging.getLogger(__name__)

class CoreVideoProcessor:
    """
    Core video processing pipeline for background replacement
    """
    
    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 = {}
        
        # 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
        }
        
        # 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}")
    
    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]:
        """
        Process video with background replacement
        
        Args:
            video_path: Input video path
            background_choice: Background type or name
            custom_background_path: Path to custom background (if applicable)
            progress_callback: Progress update callback
            cancel_event: Event to cancel processing
            preview_mask: Generate mask preview instead of final output
            preview_greenscreen: Generate greenscreen preview
            
        Returns:
            Tuple of (output_path, status_message)
        """
        if self.processing_active:
            return None, "Processing already in progress"
        
        self.processing_active = True
        start_time = time.time()
        
        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"
            
            # Process video frames
            result = self._process_video_frames(
                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"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 _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(
        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]:
        """Process all video frames"""
        
        # 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 mask cache
        self.last_refined_mask = None
        self.frame_cache.clear()
        
        try:
            prog_tracker.set_stage("Processing frames")
            
            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")
                    
                    # Process frame
                    processed_frame = self._process_single_frame(
                        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(
        self,
        frame: np.ndarray,
        background: np.ndarray,
        frame_number: int,
        preview_mask: bool,
        preview_greenscreen: bool
    ) -> np.ndarray:
        """Process a single video frame"""
        
        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(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(mask, frame_number)
            
            # Generate output based on mode
            if preview_mask:
                return self._create_mask_preview(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")
            
            return mask
            
        except Exception as e:
            self.stats['segmentation_errors'] += 1
            raise exceptions.SegmentationError(frame_number, f"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(self, frame: np.ndarray, mask: np.ndarray, frame_number: int) -> np.ndarray:
        """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:
        """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 _apply_temporal_consistency(self, current_mask: np.ndarray, frame_number: int) -> np.ndarray:
        """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(self, frame: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """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
        
        Args:
            background_choice: Background type or name
            custom_background_path: Path to custom background
            width: Target width
            height: Target height
            
        Returns:
            Prepared background image or None if failed
        """
        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"""
        return {
            '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
        }
    
    def get_status(self) -> Dict[str, Any]:
        """Get current processor status"""
        return {
            '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()
        }
    
    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
        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}")