media-engine 0.1.0__py3-none-any.whl

media_engine/extractors/ocr.py

@@ -0,0 +1,268 @@
+"""OCR extraction using EasyOCR with fast MSER pre-filtering."""
+
+from __future__ import annotations
+
+import gc
+import logging
+from typing import Any
+
+import cv2
+import numpy as np
+
+from media_engine.config import DeviceType, get_device, get_settings
+from media_engine.extractors.frame_buffer import SharedFrameBuffer
+from media_engine.schemas import BoundingBox, OcrDetection, OcrResult
+
+logger = logging.getLogger(__name__)
+
+# Singleton reader instance (lazy loaded)
+_ocr_reader: Any = None
+_ocr_languages: list[str] | None = None
+
+# MSER detector (reusable, no state)
+_mser_detector: Any = None
+
+
+def _get_mser_detector() -> Any:
+    """Get or create MSER detector (singleton)."""
+    global _mser_detector
+    if _mser_detector is None:
+        _mser_detector = cv2.MSER_create(  # type: ignore[attr-defined]
+            delta=5,  # Stability threshold
+            min_area=50,  # Min region size
+            max_area=14400,  # Max region size (120x120)
+            max_variation=0.25,
+        )
+    return _mser_detector
+
+
+def has_text_regions(
+    frame: np.ndarray,
+    min_regions: int = 3,
+    min_aspect_ratio: float = 0.2,
+    max_aspect_ratio: float = 15.0,
+) -> bool:
+    """Fast detection of potential text regions using MSER.
+
+    MSER (Maximally Stable Extremal Regions) is a classic computer vision
+    algorithm that finds stable regions in images. Text characters are
+    typically stable regions with specific aspect ratios.
+
+    This is ~100x faster than deep learning OCR and can be used to skip
+    frames that definitely don't contain text.
+
+    Args:
+        frame: BGR image as numpy array
+        min_regions: Minimum text-like regions to consider "has text"
+        min_aspect_ratio: Minimum width/height ratio for text-like regions
+        max_aspect_ratio: Maximum width/height ratio for text-like regions
+
+    Returns:
+        True if frame likely contains text, False otherwise
+    """
+    # Convert to grayscale
+    if len(frame.shape) == 3:
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    else:
+        gray = frame
+
+    # Detect MSER regions
+    mser = _get_mser_detector()
+    regions, _ = mser.detectRegions(gray)
+
+    if len(regions) < min_regions:
+        return False
+
+    # Filter regions by text-like characteristics
+    text_like_count = 0
+    for region in regions:
+        # Get bounding box
+        _, _, w, h = cv2.boundingRect(region)  # type: ignore[call-overload]
+
+        if h == 0:
+            continue
+
+        aspect_ratio = w / h
+
+        # Text characters typically have aspect ratios between 0.2 and 15
+        # (narrow letters like 'i' to wide text blocks)
+        if min_aspect_ratio <= aspect_ratio <= max_aspect_ratio:
+            # Additional filter: text regions tend to be small-medium sized
+            area = w * h
+            if 100 <= area <= 50000:
+                text_like_count += 1
+
+        # Early exit if we've found enough
+        if text_like_count >= min_regions:
+            return True
+
+    return text_like_count >= min_regions
+
+
+def unload_ocr_model() -> None:
+    """Unload the EasyOCR model to free memory."""
+    global _ocr_reader, _ocr_languages
+
+    if _ocr_reader is None:
+        return
+
+    logger.info("Unloading EasyOCR model to free memory")
+
+    try:
+        import torch
+
+        del _ocr_reader
+        _ocr_reader = None
+        _ocr_languages = None
+
+        gc.collect()
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            if hasattr(torch.mps, "synchronize"):
+                torch.mps.synchronize()
+            if hasattr(torch.mps, "empty_cache"):
+                torch.mps.empty_cache()
+
+        gc.collect()
+        logger.info("EasyOCR model unloaded")
+    except Exception as e:
+        logger.warning(f"Error unloading EasyOCR model: {e}")
+        _ocr_reader = None
+        _ocr_languages = None
+
+
+def _get_ocr_reader(languages: list[str] | None = None) -> Any:
+    """Get or create the EasyOCR reader (singleton).
+
+    Note: Once initialized, the reader keeps its languages.
+    To change languages, restart the server.
+    """
+    global _ocr_reader, _ocr_languages
+
+    if _ocr_reader is not None:
+        return _ocr_reader
+
+    import easyocr  # type: ignore[import-not-found]
+
+    if languages is None:
+        # Get from settings
+        settings = get_settings()
+        languages = settings.ocr_languages
+
+    _ocr_languages = languages
+
+    # Enable GPU if available (CUDA only - EasyOCR doesn't support MPS)
+    device = get_device()
+    use_gpu = device == DeviceType.CUDA
+    device_name = "CUDA GPU" if use_gpu else "CPU"
+
+    logger.info(f"Initializing EasyOCR with languages: {languages} on {device_name}")
+    _ocr_reader = easyocr.Reader(languages, gpu=use_gpu)
+
+    return _ocr_reader
+
+
+def extract_ocr(
+    file_path: str,
+    frame_buffer: SharedFrameBuffer,
+    min_confidence: float = 0.5,
+    languages: list[str] | None = None,
+    skip_prefilter: bool = False,
+) -> OcrResult:
+    """Extract text from video frames using two-phase OCR.
+
+    Phase 1: Fast MSER-based text detection (~5ms/frame)
+    Phase 2: Deep learning OCR on frames with text (~500ms/frame)
+
+    This typically skips 80-90% of frames, providing major speedup.
+
+    Args:
+        file_path: Path to video file (used for logging)
+        frame_buffer: Pre-decoded frames from SharedFrameBuffer
+        min_confidence: Minimum detection confidence
+        languages: OCR languages (default from settings)
+        skip_prefilter: If True, skip MSER pre-filter and run OCR on all frames
+
+    Returns:
+        OcrResult with detected text
+    """
+    detections: list[OcrDetection] = []
+    seen_texts: set[str] = set()  # For deduplication
+
+    # Stats for logging
+    frames_checked = 0
+    frames_with_text = 0
+    frames_skipped = 0
+
+    # Lazy-load OCR reader only if we find frames with text
+    reader: Any = None
+
+    def process_frame(frame: np.ndarray, timestamp: float) -> None:
+        """Process a single frame for OCR."""
+        nonlocal frames_checked, frames_with_text, frames_skipped, reader
+
+        frames_checked += 1
+
+        # Phase 1: Fast MSER pre-filter
+        if not skip_prefilter:
+            if not has_text_regions(frame):
+                frames_skipped += 1
+                return
+
+        frames_with_text += 1
+
+        # Phase 2: Deep learning OCR (only on frames that passed pre-filter)
+        try:
+            # Lazy load OCR reader on first use
+            if reader is None:
+                reader = _get_ocr_reader(languages)
+
+            results = reader.readtext(frame)
+
+            for bbox_points, text, confidence in results:
+                if confidence < min_confidence:
+                    continue
+
+                # Skip if we've seen this exact text recently
+                text_key = text.strip().lower()
+                if text_key in seen_texts:
+                    continue
+                seen_texts.add(text_key)
+
+                # Convert polygon to bounding box
+                # bbox_points is [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
+                x_coords = [p[0] for p in bbox_points]
+                y_coords = [p[1] for p in bbox_points]
+                x = int(min(x_coords))
+                y = int(min(y_coords))
+                width = int(max(x_coords) - x)
+                height = int(max(y_coords) - y)
+
+                detection = OcrDetection(
+                    timestamp=timestamp,
+                    text=text.strip(),
+                    confidence=round(float(confidence), 3),
+                    bbox=BoundingBox(x=x, y=y, width=width, height=height),
+                )
+                detections.append(detection)
+
+        except Exception as e:
+            logger.warning(f"Failed to process frame at {timestamp}s: {e}")
+
+    # Process frames from shared buffer
+    logger.info(f"Processing {len(frame_buffer.frames)} frames for OCR")
+    for ts in sorted(frame_buffer.frames.keys()):
+        shared_frame = frame_buffer.frames[ts]
+        process_frame(shared_frame.bgr, ts)
+
+    # Log stats
+    if frames_checked > 0:
+        skip_pct = (frames_skipped / frames_checked) * 100
+        logger.info(f"OCR: {frames_checked} frames checked, {frames_skipped} skipped ({skip_pct:.0f}%), " f"{frames_with_text} processed, {len(detections)} text regions found")
+    else:
+        logger.info("OCR: no frames to process")
+
+    return OcrResult(detections=detections)

media_engine/extractors/scenes.py

@@ -0,0 +1,82 @@
+"""Scene detection using PySceneDetect."""
+
+import logging
+from pathlib import Path
+
+from media_engine.schemas import SceneDetection, ScenesResult
+
+logger = logging.getLogger(__name__)
+
+# Resolution thresholds for frame skipping (pixel count)
+RES_1080P = 1920 * 1080  # ~2M pixels
+RES_4K = 3840 * 2160  # ~8.3M pixels
+RES_5K = 5120 * 2880  # ~14.7M pixels
+
+
+def extract_scenes(file_path: str, threshold: float = 27.0) -> ScenesResult:
+    """Detect scene boundaries in video file.
+
+    For high-resolution videos (4K+), uses frame skipping to improve performance.
+
+    Args:
+        file_path: Path to video file
+        threshold: Content detector threshold (lower = more sensitive)
+
+    Returns:
+        ScenesResult with detected scene boundaries
+    """
+    from scenedetect import (  # type: ignore[import-not-found]
+        ContentDetector,
+        SceneManager,
+        open_video,
+    )
+
+    path = Path(file_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Video file not found: {file_path}")
+
+    logger.info(f"Detecting scenes in {file_path}")
+
+    # Open video to get resolution
+    video = open_video(file_path)
+    width, height = video.frame_size
+    pixels = width * height
+
+    # Determine frame skip based on resolution
+    # Higher resolution = more frame skip for speed
+    if pixels > RES_5K:
+        frame_skip = 4  # Process every 5th frame for 5K+
+        logger.info(f"High-res video ({width}x{height}), using frame_skip=4")
+    elif pixels > RES_4K:
+        frame_skip = 2  # Process every 3rd frame for 4K+
+        logger.info(f"4K video ({width}x{height}), using frame_skip=2")
+    elif pixels > RES_1080P:
+        frame_skip = 1  # Process every 2nd frame for >1080p
+        logger.info(f"High-res video ({width}x{height}), using frame_skip=1")
+    else:
+        frame_skip = 0  # Process every frame for 1080p and below
+
+    # Use SceneManager API for frame_skip support
+    scene_manager = SceneManager()
+    scene_manager.add_detector(ContentDetector(threshold=threshold))
+
+    # Detect scenes with frame skipping
+    scene_manager.detect_scenes(video, frame_skip=frame_skip)
+    scenes = scene_manager.get_scene_list()
+
+    detections = []
+    for i, (start_time, end_time) in enumerate(scenes):
+        start_sec = start_time.get_seconds()
+        end_sec = end_time.get_seconds()
+        detections.append(
+            SceneDetection(
+                index=i,
+                start=start_sec,
+                end=end_sec,
+                duration=round(end_sec - start_sec, 3),
+            )
+        )
+
+    logger.info(f"Detected {len(detections)} scenes")
+
+    return ScenesResult(count=len(detections), detections=detections)

media_engine/extractors/shot_type.py

@@ -0,0 +1,217 @@
+"""Shot type detection using CLIP classification."""
+
+import logging
+import os
+import tempfile
+from enum import StrEnum
+from pathlib import Path
+from typing import Final
+
+import cv2
+
+from media_engine.config import has_cuda, is_apple_silicon
+from media_engine.extractors.frames import FrameExtractor
+from media_engine.schemas import ShotType
+
+logger = logging.getLogger(__name__)
+
+
+class ShotTypeLabel(StrEnum):
+    """Shot type classification labels."""
+
+    AERIAL = "aerial"
+    INTERVIEW = "interview"
+    B_ROLL = "b-roll"
+    STUDIO = "studio"
+    HANDHELD = "handheld"
+    STATIC = "static"
+    PHONE = "phone"
+    DASHCAM = "dashcam"
+    SECURITY = "security"
+    BROADCAST = "broadcast"
+    UNKNOWN = "unknown"
+
+
+# Shot type labels for CLIP classification
+SHOT_TYPE_LABELS: Final[list[str]] = [
+    "aerial drone footage from above",
+    "interview with person talking to camera",
+    "b-roll footage of scenery or environment",
+    "studio footage with controlled lighting",
+    "handheld camera footage",
+    "tripod static shot",
+    "phone footage vertical video",
+    "dashcam footage from car",
+    "security camera footage",
+    "news broadcast footage",
+]
+
+# Map CLIP labels to simplified shot types
+LABEL_TO_TYPE: Final[dict[str, ShotTypeLabel]] = {
+    "aerial drone footage from above": ShotTypeLabel.AERIAL,
+    "interview with person talking to camera": ShotTypeLabel.INTERVIEW,
+    "b-roll footage of scenery or environment": ShotTypeLabel.B_ROLL,
+    "studio footage with controlled lighting": ShotTypeLabel.STUDIO,
+    "handheld camera footage": ShotTypeLabel.HANDHELD,
+    "tripod static shot": ShotTypeLabel.STATIC,
+    "phone footage vertical video": ShotTypeLabel.PHONE,
+    "dashcam footage from car": ShotTypeLabel.DASHCAM,
+    "security camera footage": ShotTypeLabel.SECURITY,
+    "news broadcast footage": ShotTypeLabel.BROADCAST,
+}
+
+
+def detect_shot_type(file_path: str, sample_count: int = 5) -> ShotType | None:
+    """Detect shot type using CLIP classification.
+
+    Args:
+        file_path: Path to video file
+        sample_count: Number of frames to sample for classification
+
+    Returns:
+        ShotType with primary type and confidence, or None if detection fails
+    """
+    path = Path(file_path)
+    if not path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    try:
+        # Get video duration (0 for images)
+        duration = _get_video_duration(file_path)
+
+        # Sample frames at regular intervals (or single frame for images)
+        temp_dir = tempfile.mkdtemp(prefix="polybos_shot_")
+
+        try:
+            frames: list[str] = []
+            if duration == 0:
+                # Image or zero-duration: use single frame at timestamp 0
+                frame_path = _extract_frame_at(file_path, temp_dir, 0.0)
+                if frame_path:
+                    frames.append(frame_path)
+            else:
+                for i in range(sample_count):
+                    timestamp = (i + 0.5) * duration / sample_count
+                    frame_path = _extract_frame_at(file_path, temp_dir, timestamp)
+                    if frame_path:
+                        frames.append(frame_path)
+
+            if not frames:
+                return None
+
+            # Classify frames using CLIP
+            votes = _classify_frames(frames)
+
+            if not votes:
+                return None
+
+            # Get most common classification
+            best_label = max(votes, key=lambda k: votes.get(k, 0))
+            confidence = votes[best_label] / len(frames)
+
+            return ShotType(
+                primary=LABEL_TO_TYPE.get(best_label, ShotTypeLabel.UNKNOWN),
+                confidence=round(confidence, 3),
+                detection_method="clip",
+            )
+
+        finally:
+            # Clean up
+            import shutil
+
+            shutil.rmtree(temp_dir, ignore_errors=True)
+
+    except Exception as e:
+        logger.warning(f"Shot type detection failed: {e}")
+        return None
+
+
+def _classify_frames(frame_paths: list[str]) -> dict[str, int]:
+    """Classify frames using CLIP and return vote counts."""
+    votes: dict[str, int] = {}
+
+    if is_apple_silicon():
+        votes = _classify_with_mlx(frame_paths)
+    else:
+        votes = _classify_with_openclip(frame_paths)
+
+    return votes
+
+
+def _classify_with_openclip(frame_paths: list[str]) -> dict[str, int]:
+    """Classify frames using OpenCLIP."""
+    import open_clip  # type: ignore[import-not-found]
+    import torch
+    from PIL import Image
+
+    device: str = "cuda" if has_cuda() else "cpu"
+
+    model, _, preprocess = open_clip.create_model_and_transforms("ViT-B-32", pretrained="openai")
+    model = model.to(device)
+    model.eval()
+
+    tokenizer = open_clip.get_tokenizer("ViT-B-32")
+    text_tokens = tokenizer(SHOT_TYPE_LABELS).to(device)
+
+    with torch.no_grad():
+        text_features = model.encode_text(text_tokens)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+    votes: dict[str, int] = {}
+
+    for frame_path in frame_paths:
+        try:
+            image = Image.open(frame_path).convert("RGB")
+            image_tensor = preprocess(image).unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                image_features = model.encode_image(image_tensor)
+                image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+
+                similarity = (image_features @ text_features.T).softmax(dim=-1)
+                best_idx = similarity.argmax().item()
+                best_label = SHOT_TYPE_LABELS[best_idx]
+
+                votes[best_label] = votes.get(best_label, 0) + 1
+
+        except Exception as e:
+            logger.warning(f"Failed to classify frame {frame_path}: {e}")
+
+    return votes
+
+
+def _classify_with_mlx(frame_paths: list[str]) -> dict[str, int]:
+    """Classify frames using MLX-CLIP (Apple Silicon)."""
+    # Fall back to OpenCLIP for now as MLX-CLIP API may vary
+    # TODO: Implement native MLX-CLIP classification
+    return _classify_with_openclip(frame_paths)
+
+
+def _extract_frame_at(file_path: str, output_dir: str, timestamp: float) -> str | None:
+    """Extract a single frame at specified timestamp.
+
+    Uses FrameExtractor which handles both videos (via OpenCV/ffmpeg)
+    and images (via direct loading).
+    """
+    output_path = os.path.join(output_dir, f"frame_{timestamp:.3f}.jpg")
+
+    with FrameExtractor(file_path) as extractor:
+        frame = extractor.get_frame_at(timestamp)
+
+        if frame is not None:
+            cv2.imwrite(output_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
+            if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
+                return output_path
+            else:
+                logger.warning(f"Frame at {timestamp}s: could not save to {output_path}")
+        else:
+            logger.warning(f"Frame at {timestamp}s: extraction failed")
+
+    return None
+
+
+def _get_video_duration(file_path: str) -> float:
+    """Get video/image duration in seconds (0 for images)."""
+    from media_engine.extractors.frames import get_video_duration
+
+    return get_video_duration(file_path)