AI-Talent-Force/dev_caio

1"""2ShortSmith v2 - Motion Detector Module3 4Motion analysis using optical flow for:5- Detecting action-heavy segments6- Identifying camera movement vs subject movement7- Dynamic FPS scaling based on motion level8 9Uses RAFT (Recurrent All-Pairs Field Transforms) for high-quality10optical flow, with fallback to Farneback for speed.11"""12 13from pathlib import Path14from typing import List, Optional, Tuple, Union15from dataclasses import dataclass16import numpy as np17 18from utils.logger import get_logger, LogTimer19from utils.helpers import ModelLoadError, InferenceError20from config import get_config, ModelConfig21 22logger = get_logger("models.motion_detector")23 24 25@dataclass26class MotionScore:27    """Motion analysis result for a frame pair."""28    timestamp: float           # Timestamp of second frame29    magnitude: float           # Average motion magnitude (0-1 normalized)30    direction: float           # Dominant motion direction (radians)31    uniformity: float          # How uniform the motion is (1 = all same direction)32    is_camera_motion: bool     # Likely camera motion vs subject motion33 34    @property35    def is_high_motion(self) -> bool:36        """Check if this is a high-motion segment."""37        return self.magnitude > 0.338 39    @property40    def is_action(self) -> bool:41        """Check if this likely contains action (non-uniform motion)."""42        return self.magnitude > 0.2 and self.uniformity < 0.743 44 45class MotionDetector:46    """47    Motion detection using optical flow.48 49    Supports:50    - RAFT optical flow (high quality, GPU)51    - Farneback optical flow (faster, CPU)52    - Motion magnitude scoring53    - Camera vs subject motion detection54    """55 56    def __init__(57        self,58        config: Optional[ModelConfig] = None,59        use_raft: bool = True,60    ):61        """62        Initialize motion detector.63 64        Args:65            config: Model configuration66            use_raft: Whether to use RAFT (True) or Farneback (False)67        """68        self.config = config or get_config().model69        self.use_raft = use_raft70        self.raft_model = None71 72        if use_raft:73            self._load_raft()74 75        logger.info(f"MotionDetector initialized (RAFT={use_raft})")76 77    def _load_raft(self) -> None:78        """Load RAFT optical flow model."""79        try:80            import torch81            from torchvision.models.optical_flow import raft_small, Raft_Small_Weights82 83            logger.info("Loading RAFT optical flow model...")84 85            weights = Raft_Small_Weights.DEFAULT86            self.raft_model = raft_small(weights=weights)87 88            if self.config.device == "cuda" and torch.cuda.is_available():89                self.raft_model = self.raft_model.cuda()90 91            self.raft_model.eval()92 93            # Store preprocessing transforms94            self._raft_transforms = weights.transforms()95 96            logger.info("RAFT model loaded successfully")97 98        except Exception as e:99            logger.warning(f"Failed to load RAFT model, using Farneback: {e}")100            self.use_raft = False101            self.raft_model = None102 103    def compute_flow(104        self,105        frame1: np.ndarray,106        frame2: np.ndarray,107    ) -> np.ndarray:108        """109        Compute optical flow between two frames.110 111        Args:112            frame1: First frame (BGR or RGB, HxWxC)113            frame2: Second frame (BGR or RGB, HxWxC)114 115        Returns:116            Optical flow array (HxWx2), flow[y,x] = (dx, dy)117        """118        if self.use_raft and self.raft_model is not None:119            return self._compute_raft_flow(frame1, frame2)120        else:121            return self._compute_farneback_flow(frame1, frame2)122 123    def _compute_raft_flow(124        self,125        frame1: np.ndarray,126        frame2: np.ndarray,127    ) -> np.ndarray:128        """Compute flow using RAFT."""129        import torch130 131        try:132            # Convert to RGB if BGR133            if frame1.shape[2] == 3:134                frame1_rgb = frame1[:, :, ::-1].copy()135                frame2_rgb = frame2[:, :, ::-1].copy()136            else:137                frame1_rgb = frame1138                frame2_rgb = frame2139 140            # Convert to tensors141            img1 = torch.from_numpy(frame1_rgb).permute(2, 0, 1).float().unsqueeze(0)142            img2 = torch.from_numpy(frame2_rgb).permute(2, 0, 1).float().unsqueeze(0)143 144            if self.config.device == "cuda" and torch.cuda.is_available():145                img1 = img1.cuda()146                img2 = img2.cuda()147 148            # Compute flow149            with torch.no_grad():150                flow_predictions = self.raft_model(img1, img2)151                flow = flow_predictions[-1]  # Use final prediction152 153            # Convert back to numpy154            flow = flow[0].permute(1, 2, 0).cpu().numpy()155 156            return flow157 158        except Exception as e:159            logger.warning(f"RAFT flow failed, using Farneback: {e}")160            return self._compute_farneback_flow(frame1, frame2)161 162    def _compute_farneback_flow(163        self,164        frame1: np.ndarray,165        frame2: np.ndarray,166    ) -> np.ndarray:167        """Compute flow using Farneback algorithm."""168        import cv2169 170        # Convert to grayscale171        if len(frame1.shape) == 3:172            gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)173            gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)174        else:175            gray1 = frame1176            gray2 = frame2177 178        # Compute Farneback optical flow179        flow = cv2.calcOpticalFlowFarneback(180            gray1, gray2,181            None,182            pyr_scale=0.5,183            levels=3,184            winsize=15,185            iterations=3,186            poly_n=5,187            poly_sigma=1.2,188            flags=0,189        )190 191        return flow192 193    def analyze_motion(194        self,195        frame1: np.ndarray,196        frame2: np.ndarray,197        timestamp: float = 0.0,198    ) -> MotionScore:199        """200        Analyze motion between two frames.201 202        Args:203            frame1: First frame204            frame2: Second frame205            timestamp: Timestamp of second frame206 207        Returns:208            MotionScore with analysis results209        """210        flow = self.compute_flow(frame1, frame2)211 212        # Compute magnitude and direction213        magnitude = np.sqrt(flow[:, :, 0]**2 + flow[:, :, 1]**2)214        direction = np.arctan2(flow[:, :, 1], flow[:, :, 0])215 216        # Average magnitude (normalized by image diagonal)217        h, w = frame1.shape[:2]218        diagonal = np.sqrt(h**2 + w**2)219        avg_magnitude = float(np.mean(magnitude) / diagonal)220 221        # Dominant direction222        # Weight by magnitude to get dominant direction223        weighted_direction = np.average(direction, weights=magnitude + 1e-8)224 225        # Uniformity: how consistent is the motion direction?226        # High uniformity = likely camera motion227        dir_std = float(np.std(direction))228        uniformity = 1.0 / (1.0 + dir_std)229 230        # Detect camera motion (uniform direction across frame)231        is_camera = uniformity > 0.7 and avg_magnitude > 0.05232 233        return MotionScore(234            timestamp=timestamp,235            magnitude=min(1.0, avg_magnitude * 10),  # Scale up236            direction=float(weighted_direction),237            uniformity=uniformity,238            is_camera_motion=is_camera,239        )240 241    def analyze_video_segment(242        self,243        frames: List[np.ndarray],244        timestamps: List[float],245    ) -> List[MotionScore]:246        """247        Analyze motion across a video segment.248 249        Args:250            frames: List of frames251            timestamps: Timestamps for each frame252 253        Returns:254            List of MotionScore objects (one per frame pair)255        """256        if len(frames) < 2:257            return []258 259        scores = []260 261        with LogTimer(logger, f"Analyzing motion in {len(frames)} frames"):262            for i in range(1, len(frames)):263                try:264                    score = self.analyze_motion(265                        frames[i-1],266                        frames[i],267                        timestamps[i],268                    )269                    scores.append(score)270                except Exception as e:271                    logger.warning(f"Motion analysis failed for frame {i}: {e}")272 273        return scores274 275    def get_motion_heatmap(276        self,277        frame1: np.ndarray,278        frame2: np.ndarray,279    ) -> np.ndarray:280        """281        Get motion magnitude heatmap.282 283        Args:284            frame1: First frame285            frame2: Second frame286 287        Returns:288            Heatmap of motion magnitude (HxW, values 0-255)289        """290        flow = self.compute_flow(frame1, frame2)291        magnitude = np.sqrt(flow[:, :, 0]**2 + flow[:, :, 1]**2)292 293        # Normalize to 0-255294        max_mag = np.percentile(magnitude, 99)  # Robust max295        if max_mag > 0:296            normalized = np.clip(magnitude / max_mag * 255, 0, 255)297        else:298            normalized = np.zeros_like(magnitude)299 300        return normalized.astype(np.uint8)301 302    def compute_aggregate_motion(303        self,304        scores: List[MotionScore],305    ) -> float:306        """307        Compute aggregate motion score for a segment.308 309        Args:310            scores: List of MotionScore objects311 312        Returns:313            Aggregate motion score (0-1)314        """315        if not scores:316            return 0.0317 318        # Weight by non-camera motion319        weighted_sum = sum(320            s.magnitude * (0.3 if s.is_camera_motion else 1.0)321            for s in scores322        )323 324        return weighted_sum / len(scores)325 326    def identify_high_motion_segments(327        self,328        scores: List[MotionScore],329        threshold: float = 0.3,330        min_duration: int = 3,331    ) -> List[Tuple[float, float, float]]:332        """333        Identify segments with high motion.334 335        Args:336            scores: List of MotionScore objects337            threshold: Minimum motion magnitude338            min_duration: Minimum number of consecutive frames339 340        Returns:341            List of (start_time, end_time, avg_motion) tuples342        """343        if not scores:344            return []345 346        segments = []347        in_segment = False348        segment_start = 0.0349        segment_scores = []350 351        for score in scores:352            if score.magnitude >= threshold:353                if not in_segment:354                    in_segment = True355                    segment_start = score.timestamp356                    segment_scores = [score.magnitude]357                else:358                    segment_scores.append(score.magnitude)359            else:360                if in_segment:361                    if len(segment_scores) >= min_duration:362                        segments.append((363                            segment_start,364                            score.timestamp,365                            sum(segment_scores) / len(segment_scores),366                        ))367                    in_segment = False368 369        # Handle segment at end370        if in_segment and len(segment_scores) >= min_duration:371            segments.append((372                segment_start,373                scores[-1].timestamp,374                sum(segment_scores) / len(segment_scores),375            ))376 377        logger.info(f"Found {len(segments)} high-motion segments")378        return segments379 380 381# Export public interface382__all__ = ["MotionDetector", "MotionScore"]383