Update pipeline/orchestrator.py

pipeline/orchestrator.py

@@ -14,7 +14,7 @@ Manages the flow between all components:
 """
 
 from pathlib import Path
-from typing import List, Optional, Callable, Dict, Any, Generator
+from typing import List, Optional, Callable, Dict, Any
 from dataclasses import dataclass, field
 from enum import Enum
 import time
@@ -39,7 +39,7 @@ from models.face_recognizer import FaceRecognizer
 from models.body_recognizer import BodyRecognizer
 from models.motion_detector import MotionDetector
 from scoring.hype_scorer import HypeScorer, SegmentScore
-from scoring.domain_presets import get_domain_preset, Domain
+from scoring.domain_presets import get_domain_preset
 from scoring.viral_hooks import ViralHookDetector, HookSignal
 
 logger = get_logger("pipeline.orchestrator")
@@ -294,6 +294,7 @@ class PipelineOrchestrator:
 
             # Visual analysis (if enabled)
             visual_features = []
+            custom_analysis_results = []
             if self._visual_analyzer is not None:
                 self._update_progress(PipelineStage.ANALYZING_VISUAL, 0.0, "Analyzing visual content...")
                 try:
@@ -302,11 +303,26 @@ class PipelineOrchestrator:
                             frame.frame_path, timestamp=frame.timestamp
                         )
                         visual_features.append(features)
+
+                        # Apply custom prompt analysis if provided
+                        if custom_prompt:
+                            custom_result = self._visual_analyzer.analyze_with_custom_prompt(
+                                frame.frame_path,
+                                prompt=custom_prompt,
+                                timestamp=frame.timestamp,
+                            )
+                            custom_analysis_results.append(custom_result)
+
                         self._update_progress(
                             PipelineStage.ANALYZING_VISUAL,
                             (i + 1) / len(frames),
                             f"Analyzing frame {i+1}/{len(frames)}"
                         )
+
+                    # Boost scores based on custom prompt matches
+                    if custom_analysis_results:
+                        self._apply_custom_prompt_boost(visual_features, custom_analysis_results)
+
                 except Exception as e:
                     logger.warning(f"Visual analysis failed, continuing without: {e}")
             self._update_progress(PipelineStage.ANALYZING_VISUAL, 1.0, "Visual analysis complete")
@@ -346,9 +362,12 @@ class PipelineOrchestrator:
                     logger.warning(f"Person detection failed: {e}")
             self._update_progress(PipelineStage.DETECTING_PERSON, 1.0, "Person detection complete")
 
-            # Motion analysis (simplified)
+            # Motion analysis
             self._update_progress(PipelineStage.ANALYZING_MOTION, 0.0, "Analyzing motion...")
-            motion_scores = self._estimate_motion_from_visual(visual_features)
+            motion_scores = self._compute_motion_scores(frames)
+            # Fallback to visual estimation if motion detector failed or unavailable
+            if not motion_scores and visual_features:
+                motion_scores = self._estimate_motion_from_visual(visual_features)
             self._update_progress(PipelineStage.ANALYZING_MOTION, 1.0, "Motion analysis complete")
 
             # Scoring
@@ -466,6 +485,16 @@ class PipelineOrchestrator:
             logger.warning(f"Visual analyzer not available: {e}")
             self._visual_analyzer = None
 
+        # Motion detector (optional, falls back to visual estimation)
+        try:
+            self._motion_detector = MotionDetector(
+                self.config.model,
+                use_raft=True,  # Use high-quality RAFT if available
+            )
+        except Exception as e:
+            logger.warning(f"Motion detector not available, using visual estimation: {e}")
+            self._motion_detector = None
+
         # Person recognition (only if needed)
         if person_filter:
             try:
@@ -586,11 +615,62 @@ class PipelineOrchestrator:
             for s in scores
         ]
 
+    def _compute_motion_scores(
+        self,
+        frames: List[SampledFrame],
+    ) -> List[float]:
+        """
+        Compute motion scores using MotionDetector or fallback to visual estimation.
+
+        Args:
+            frames: Sampled frames with paths and timestamps
+
+        Returns:
+            List of motion scores (0-1) for each frame
+        """
+        if not frames:
+            return []
+
+        # Use real motion detector if available
+        if self._motion_detector is not None and len(frames) >= 2:
+            try:
+                import cv2
+
+                motion_scores = []
+
+                # Load frames and compute motion between consecutive pairs
+                prev_frame = None
+                for i, frame in enumerate(frames):
+                    curr_frame = cv2.imread(str(frame.frame_path))
+
+                    if prev_frame is not None and curr_frame is not None:
+                        motion_result = self._motion_detector.analyze_motion(
+                            prev_frame, curr_frame, timestamp=frame.timestamp
+                        )
+                        motion_scores.append(motion_result.magnitude)
+                    else:
+                        # First frame has no motion score
+                        if i == 0:
+                            motion_scores.append(0.0)
+
+                    prev_frame = curr_frame
+
+                logger.info(f"Computed motion scores for {len(motion_scores)} frames using RAFT/Farneback")
+                return motion_scores
+
+            except Exception as e:
+                logger.warning(f"Motion detection failed, falling back to visual estimation: {e}")
+
+        # Fallback: estimate from visual features (requires visual_features from caller)
+        # Return empty list - will be filled by visual estimation in scoring
+        logger.info("Using visual estimation for motion scores")
+        return []
+
     def _estimate_motion_from_visual(
         self,
         visual_features: List[VisualFeatures],
     ) -> List[float]:
-        """Estimate motion scores from visual analysis."""
+        """Estimate motion scores from visual analysis (fallback)."""
         if not visual_features:
             return []
 
@@ -608,6 +688,53 @@ class PipelineOrchestrator:
 
         return [motion_map.get(f.action_detected, 0.4) for f in visual_features]
 
+    def _apply_custom_prompt_boost(
+        self,
+        visual_features: List[VisualFeatures],
+        custom_results: List[Dict],
+    ) -> None:
+        """
+        Boost visual scores based on custom prompt responses.
+
+        Analyzes custom prompt responses and boosts hype scores for frames
+        where the response indicates a match with the user's criteria.
+
+        Args:
+            visual_features: Visual features to modify (in-place)
+            custom_results: Results from custom prompt analysis
+        """
+        if not custom_results or len(custom_results) != len(visual_features):
+            return
+
+        # Keywords that indicate positive matches
+        positive_keywords = [
+            "yes", "true", "found", "detected", "present", "visible",
+            "showing", "contains", "includes", "displays", "features",
+            "action", "exciting", "highlight", "important", "key",
+            "peak", "climax", "intense", "dramatic", "significant",
+        ]
+
+        for i, (features, custom) in enumerate(zip(visual_features, custom_results)):
+            response = custom.get("response", "").lower()
+
+            # Check for positive indicators
+            match_score = 0.0
+            for keyword in positive_keywords:
+                if keyword in response:
+                    match_score += 0.1
+
+            # Cap the boost at 50%
+            boost = min(0.5, match_score)
+
+            if boost > 0:
+                # Boost the hype score
+                original_score = features.hype_score
+                features.hype_score = min(1.0, features.hype_score * (1 + boost))
+                logger.debug(
+                    f"Frame {i}: custom prompt boost {boost:.2f} "
+                    f"({original_score:.2f} -> {features.hype_score:.2f})"
+                )
+
     def _detect_viral_hooks(
         self,
         frames: List[SampledFrame],