AI-Talent-Force/dev_caio

1"""2ShortSmith v2 - Frame Sampler Module3 4Hierarchical frame sampling strategy:51. Coarse pass: Sample 1 frame per N seconds to identify candidate regions62. Dense pass: Sample at higher FPS only on promising segments73. Dynamic FPS: Adjust sampling based on motion/content8"""9 10from pathlib import Path11from typing import List, Optional, Tuple, Generator12from dataclasses import dataclass, field13import numpy as np14 15from utils.logger import get_logger, LogTimer16from utils.helpers import VideoProcessingError, batch_list17from config import get_config, ProcessingConfig18from core.video_processor import VideoProcessor, VideoMetadata19 20logger = get_logger("core.frame_sampler")21 22 23@dataclass24class SampledFrame:25    """Represents a sampled frame with metadata."""26    frame_path: Path          # Path to the frame image file27    timestamp: float          # Timestamp in seconds28    frame_index: int          # Index in the video29    is_dense_sample: bool     # Whether from dense sampling pass30    scene_id: Optional[int] = None  # Associated scene ID31 32    # Optional: frame data loaded into memory33    frame_data: Optional[np.ndarray] = field(default=None, repr=False)34 35    @property36    def filename(self) -> str:37        """Get the frame filename."""38        return self.frame_path.name39 40 41@dataclass42class SamplingRegion:43    """A region identified for dense sampling."""44    start_time: float45    end_time: float46    priority_score: float  # Higher = more likely to contain highlights47 48    @property49    def duration(self) -> float:50        return self.end_time - self.start_time51 52 53class FrameSampler:54    """55    Intelligent frame sampler using hierarchical strategy.56 57    Optimizes compute by:58    1. Sparse sampling to identify candidate regions59    2. Dense sampling only on promising areas60    3. Skipping static/low-motion content61    """62 63    def __init__(64        self,65        video_processor: VideoProcessor,66        config: Optional[ProcessingConfig] = None,67    ):68        """69        Initialize frame sampler.70 71        Args:72            video_processor: VideoProcessor instance for frame extraction73            config: Processing configuration (uses default if None)74        """75        self.video_processor = video_processor76        self.config = config or get_config().processing77 78        logger.info(79            f"FrameSampler initialized (coarse={self.config.coarse_sample_interval}s, "80            f"dense_fps={self.config.dense_sample_fps})"81        )82 83    def sample_coarse(84        self,85        video_path: str | Path,86        output_dir: str | Path,87        metadata: Optional[VideoMetadata] = None,88        start_time: float = 0,89        end_time: Optional[float] = None,90    ) -> List[SampledFrame]:91        """92        Perform coarse sampling pass.93 94        Samples 1 frame every N seconds (default 5s) across the video.95 96        Args:97            video_path: Path to the video file98            output_dir: Directory to save extracted frames99            metadata: Video metadata (fetched if not provided)100            start_time: Start sampling from this timestamp101            end_time: End sampling at this timestamp102 103        Returns:104            List of SampledFrame objects105        """106        video_path = Path(video_path)107        output_dir = Path(output_dir)108        output_dir.mkdir(parents=True, exist_ok=True)109 110        # Get metadata if not provided111        if metadata is None:112            metadata = self.video_processor.get_metadata(video_path)113 114        end_time = end_time or metadata.duration115 116        # Validate time range117        if end_time > metadata.duration:118            end_time = metadata.duration119        if start_time >= end_time:120            raise VideoProcessingError(121                f"Invalid time range: {start_time} to {end_time}"122            )123 124        with LogTimer(logger, f"Coarse sampling {video_path.name}"):125            # Calculate timestamps126            interval = self.config.coarse_sample_interval127            timestamps = []128            current = start_time129 130            while current < end_time:131                timestamps.append(current)132                current += interval133 134            logger.info(135                f"Coarse sampling: {len(timestamps)} frames "136                f"({interval}s interval over {end_time - start_time:.1f}s)"137            )138 139            # Extract frames140            frame_paths = self.video_processor.extract_frames(141                video_path,142                output_dir / "coarse",143                timestamps=timestamps,144            )145 146            # Create SampledFrame objects147            frames = []148            for i, (path, ts) in enumerate(zip(frame_paths, timestamps)):149                frames.append(SampledFrame(150                    frame_path=path,151                    timestamp=ts,152                    frame_index=int(ts * metadata.fps),153                    is_dense_sample=False,154                ))155 156            return frames157 158    def sample_dense(159        self,160        video_path: str | Path,161        output_dir: str | Path,162        regions: List[SamplingRegion],163        metadata: Optional[VideoMetadata] = None,164    ) -> List[SampledFrame]:165        """166        Perform dense sampling on specific regions.167 168        Args:169            video_path: Path to the video file170            output_dir: Directory to save extracted frames171            regions: List of regions to sample densely172            metadata: Video metadata (fetched if not provided)173 174        Returns:175            List of SampledFrame objects from dense regions176        """177        video_path = Path(video_path)178        output_dir = Path(output_dir)179 180        if metadata is None:181            metadata = self.video_processor.get_metadata(video_path)182 183        all_frames = []184 185        with LogTimer(logger, f"Dense sampling {len(regions)} regions"):186            for i, region in enumerate(regions):187                region_dir = output_dir / f"dense_region_{i:03d}"188                region_dir.mkdir(parents=True, exist_ok=True)189 190                logger.debug(191                    f"Dense sampling region {i}: "192                    f"{region.start_time:.1f}s - {region.end_time:.1f}s"193                )194 195                # Extract at dense FPS196                frame_paths = self.video_processor.extract_frames(197                    video_path,198                    region_dir,199                    fps=self.config.dense_sample_fps,200                    start_time=region.start_time,201                    end_time=region.end_time,202                )203 204                # Calculate timestamps for each frame205                for j, path in enumerate(frame_paths):206                    timestamp = region.start_time + (j / self.config.dense_sample_fps)207                    all_frames.append(SampledFrame(208                        frame_path=path,209                        timestamp=timestamp,210                        frame_index=int(timestamp * metadata.fps),211                        is_dense_sample=True,212                    ))213 214            logger.info(f"Dense sampling extracted {len(all_frames)} frames")215            return all_frames216 217    def sample_hierarchical(218        self,219        video_path: str | Path,220        output_dir: str | Path,221        candidate_scorer: Optional[callable] = None,222        top_k_regions: int = 5,223        metadata: Optional[VideoMetadata] = None,224    ) -> Tuple[List[SampledFrame], List[SampledFrame]]:225        """226        Perform full hierarchical sampling.227 228        1. Coarse pass to identify candidates229        2. Score candidate regions230        3. Dense pass on top-k regions231 232        Args:233            video_path: Path to the video file234            output_dir: Directory to save extracted frames235            candidate_scorer: Function to score candidate regions (optional)236            top_k_regions: Number of top regions to densely sample237            metadata: Video metadata (fetched if not provided)238 239        Returns:240            Tuple of (coarse_frames, dense_frames)241        """242        video_path = Path(video_path)243        output_dir = Path(output_dir)244 245        if metadata is None:246            metadata = self.video_processor.get_metadata(video_path)247 248        with LogTimer(logger, "Hierarchical sampling"):249            # Step 1: Coarse sampling250            coarse_frames = self.sample_coarse(251                video_path, output_dir, metadata252            )253 254            # Step 2: Identify candidate regions255            if candidate_scorer is not None:256                # Use provided scorer to identify promising regions257                regions = self._identify_candidate_regions(258                    coarse_frames, candidate_scorer, top_k_regions259                )260            else:261                # Default: uniform distribution262                regions = self._create_uniform_regions(263                    metadata.duration, top_k_regions264                )265 266            # Step 3: Dense sampling on top regions267            dense_frames = self.sample_dense(268                video_path, output_dir, regions, metadata269            )270 271            logger.info(272                f"Hierarchical sampling complete: "273                f"{len(coarse_frames)} coarse, {len(dense_frames)} dense frames"274            )275 276            return coarse_frames, dense_frames277 278    def _identify_candidate_regions(279        self,280        frames: List[SampledFrame],281        scorer: callable,282        top_k: int,283    ) -> List[SamplingRegion]:284        """285        Identify top candidate regions based on scoring.286 287        Args:288            frames: List of coarse sampled frames289            scorer: Function that takes frame and returns score (0-1)290            top_k: Number of regions to return291 292        Returns:293            List of SamplingRegion objects294        """295        # Score each frame296        scores = []297        for frame in frames:298            try:299                score = scorer(frame)300                scores.append((frame, score))301            except Exception as e:302                logger.warning(f"Failed to score frame {frame.timestamp}s: {e}")303                scores.append((frame, 0.0))304 305        # Sort by score306        scores.sort(key=lambda x: x[1], reverse=True)307 308        # Create regions around top frames309        interval = self.config.coarse_sample_interval310        regions = []311 312        for frame, score in scores[:top_k]:313            # Expand region around this frame314            start = max(0, frame.timestamp - interval)315            end = frame.timestamp + interval316 317            regions.append(SamplingRegion(318                start_time=start,319                end_time=end,320                priority_score=score,321            ))322 323        # Merge overlapping regions324        regions = self._merge_overlapping_regions(regions)325 326        return regions327 328    def _create_uniform_regions(329        self,330        duration: float,331        num_regions: int,332    ) -> List[SamplingRegion]:333        """334        Create uniformly distributed sampling regions.335 336        Args:337            duration: Total video duration338            num_regions: Number of regions to create339 340        Returns:341            List of uniformly spaced SamplingRegion objects342        """343        region_duration = self.config.coarse_sample_interval * 2344        gap = (duration - region_duration * num_regions) / (num_regions + 1)345 346        if gap < 0:347            # Video too short, create fewer regions348            gap = 0349            num_regions = max(1, int(duration / region_duration))350 351        regions = []352        current = gap353 354        for i in range(num_regions):355            regions.append(SamplingRegion(356                start_time=current,357                end_time=min(current + region_duration, duration),358                priority_score=1.0 / num_regions,359            ))360            current += region_duration + gap361 362        return regions363 364    def _merge_overlapping_regions(365        self,366        regions: List[SamplingRegion],367    ) -> List[SamplingRegion]:368        """369        Merge overlapping sampling regions.370 371        Args:372            regions: List of potentially overlapping regions373 374        Returns:375            List of merged regions376        """377        if not regions:378            return []379 380        # Sort by start time381        sorted_regions = sorted(regions, key=lambda r: r.start_time)382        merged = [sorted_regions[0]]383 384        for region in sorted_regions[1:]:385            last = merged[-1]386 387            if region.start_time <= last.end_time:388                # Merge389                merged[-1] = SamplingRegion(390                    start_time=last.start_time,391                    end_time=max(last.end_time, region.end_time),392                    priority_score=max(last.priority_score, region.priority_score),393                )394            else:395                merged.append(region)396 397        return merged398 399    def sample_at_timestamps(400        self,401        video_path: str | Path,402        output_dir: str | Path,403        timestamps: List[float],404        metadata: Optional[VideoMetadata] = None,405    ) -> List[SampledFrame]:406        """407        Sample frames at specific timestamps.408 409        Args:410            video_path: Path to the video file411            output_dir: Directory to save extracted frames412            timestamps: List of timestamps to sample413            metadata: Video metadata (fetched if not provided)414 415        Returns:416            List of SampledFrame objects417        """418        video_path = Path(video_path)419        output_dir = Path(output_dir)420        output_dir.mkdir(parents=True, exist_ok=True)421 422        if metadata is None:423            metadata = self.video_processor.get_metadata(video_path)424 425        with LogTimer(logger, f"Sampling {len(timestamps)} specific timestamps"):426            frame_paths = self.video_processor.extract_frames(427                video_path,428                output_dir / "specific",429                timestamps=timestamps,430            )431 432            frames = []433            for path, ts in zip(frame_paths, timestamps):434                frames.append(SampledFrame(435                    frame_path=path,436                    timestamp=ts,437                    frame_index=int(ts * metadata.fps),438                    is_dense_sample=False,439                ))440 441            return frames442 443    def get_keyframes(444        self,445        video_path: str | Path,446        output_dir: str | Path,447        scenes: Optional[List] = None,448    ) -> List[SampledFrame]:449        """450        Extract keyframes (one per scene).451 452        Args:453            video_path: Path to the video file454            output_dir: Directory to save extracted frames455            scenes: List of Scene objects (detected if not provided)456 457        Returns:458            List of keyframe SampledFrame objects459        """460        from core.scene_detector import SceneDetector461 462        video_path = Path(video_path)463 464        if scenes is None:465            detector = SceneDetector()466            scenes = detector.detect_scenes(video_path)467 468        # Get midpoint of each scene as keyframe469        timestamps = [scene.midpoint for scene in scenes]470 471        with LogTimer(logger, f"Extracting {len(timestamps)} keyframes"):472            frames = self.sample_at_timestamps(473                video_path, output_dir, timestamps474            )475 476            # Add scene IDs477            for frame, scene_id in zip(frames, range(len(scenes))):478                frame.scene_id = scene_id479 480            return frames481 482 483# Export public interface484__all__ = ["FrameSampler", "SampledFrame", "SamplingRegion"]485