media-engine 0.1.0__py3-none-any.whl

media_engine/extractors/frames.py

@@ -0,0 +1,402 @@
+"""Fast frame extraction using OpenCV or FFmpeg.
+
+OpenCV's VideoCapture is fast for normal videos but decodes at full resolution.
+For high-resolution videos (4K+), FFmpeg decoding at target resolution is faster.
+
+Also supports direct image loading - when given an image file, it loads it
+directly instead of trying to use VideoCapture.
+"""
+
+import logging
+import os
+import subprocess
+import tempfile
+
+import cv2
+import numpy as np
+
+from media_engine.schemas import MediaType, get_media_type
+
+logger = logging.getLogger(__name__)
+
+# Resolution threshold for using FFmpeg decode (4K+)
+HIGH_RES_THRESHOLD = 3840 * 2160  # ~8.3M pixels
+
+
+class FrameExtractor:
+    """Extract frames from video or image using OpenCV.
+
+    Uses cv2.VideoCapture for fast seeking and frame extraction from videos.
+    Falls back to ffmpeg for exotic codecs that OpenCV can't handle.
+    For images, loads directly with cv2.imread (no frame extraction needed).
+    """
+
+    # Default max dimension - scale down 4K to ~HD for faster processing
+    DEFAULT_MAX_DIMENSION = 1920
+
+    def __init__(self, file_path: str, max_dimension: int | None = DEFAULT_MAX_DIMENSION):
+        """Initialize frame extractor.
+
+        Args:
+            file_path: Path to video or image file
+            max_dimension: Maximum width/height. Frames larger than this are scaled down.
+                          Set to None to disable scaling. Default: 1920 (HD)
+        """
+        self.video_path = file_path  # Keep name for compatibility
+        self.max_dimension = max_dimension
+        self.cap: cv2.VideoCapture | None = None
+        self._duration: float | None = None
+        self._fps: float | None = None
+        self._frame_count: int | None = None
+        self._width: int | None = None
+        self._height: int | None = None
+        self._use_ffmpeg_fallback = False
+        self._use_ffmpeg_decode = False  # For high-res, decode at lower res with FFmpeg
+        # Image handling
+        self._is_image = False
+        self._image_frame: np.ndarray | None = None
+
+    def __enter__(self) -> "FrameExtractor":
+        """Open video or image file."""
+        # Check if this is an image file
+        media_type = get_media_type(self.video_path)
+        if media_type == MediaType.IMAGE:
+            self._is_image = True
+            self._duration = 0.0
+            self._fps = 1.0
+            self._frame_count = 1
+            # Load the image directly
+            self._image_frame = cv2.imread(self.video_path)
+            if self._image_frame is None:
+                logger.warning(f"Failed to load image: {self.video_path}")
+            else:
+                # Apply scaling
+                self._image_frame = self._scale_frame(self._image_frame)
+                logger.debug(f"Loaded image directly: {self.video_path}")
+            return self
+
+        # Video file - use VideoCapture
+        self.cap = cv2.VideoCapture(self.video_path)
+
+        if not self.cap.isOpened():
+            logger.warning(f"OpenCV failed to open {self.video_path}, using ffmpeg fallback")
+            self._use_ffmpeg_fallback = True
+            self.cap = None
+        else:
+            self._fps = self.cap.get(cv2.CAP_PROP_FPS)
+            self._frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            self._width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            self._height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+            if self._fps > 0 and self._frame_count > 0:
+                self._duration = self._frame_count / self._fps
+            else:
+                # Seek to estimate duration
+                self._duration = self._get_duration_ffprobe()
+
+            # Check if this is high-res video that needs FFmpeg decode
+            if self._width and self._height and self.max_dimension:
+                pixels = self._width * self._height
+                max_dim = max(self._width, self._height)
+                if pixels > HIGH_RES_THRESHOLD and max_dim > self.max_dimension:
+                    logger.info(f"High-res video ({self._width}x{self._height}), " f"using FFmpeg decode at {self.max_dimension}px")
+                    self._use_ffmpeg_decode = True
+                    # Release opencv capture - we'll use FFmpeg instead
+                    self.cap.release()
+                    self.cap = None
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # noqa: ANN001
+        """Release video file."""
+        if self.cap is not None:
+            self.cap.release()
+            self.cap = None
+        # Clear image reference
+        self._image_frame = None
+
+    @property
+    def is_image(self) -> bool:
+        """Check if this extractor is handling an image (not a video)."""
+        return self._is_image
+
+    @property
+    def duration(self) -> float:
+        """Get video duration in seconds (0 for images)."""
+        if self._duration is None:
+            self._duration = self._get_duration_ffprobe()
+        return self._duration
+
+    @property
+    def fps(self) -> float:
+        """Get video frame rate (1 for images)."""
+        if self._fps is None:
+            self._fps = 30.0  # Default fallback
+        return self._fps
+
+    def _get_duration_ffprobe(self) -> float:
+        """Get duration using ffprobe."""
+        try:
+            cmd = [
+                "ffprobe",
+                "-v",
+                "quiet",
+                "-show_entries",
+                "format=duration",
+                "-of",
+                "default=noprint_wrappers=1:nokey=1",
+                self.video_path,
+            ]
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            return float(result.stdout.strip())
+        except (subprocess.CalledProcessError, ValueError):
+            return 0.0
+
+    def _scale_frame(self, frame: np.ndarray) -> np.ndarray:
+        """Scale down frame if larger than max_dimension.
+
+        Maintains aspect ratio. Only scales down, never up.
+        """
+        if self.max_dimension is None:
+            return frame
+
+        h, w = frame.shape[:2]
+        max_dim = max(h, w)
+
+        if max_dim <= self.max_dimension:
+            return frame
+
+        # Calculate scale factor
+        scale = self.max_dimension / max_dim
+        new_w = int(w * scale)
+        new_h = int(h * scale)
+
+        # Use INTER_AREA for downscaling (best quality)
+        return cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
+    def get_frame_at(self, timestamp: float) -> np.ndarray | None:
+        """Extract a single frame at the given timestamp.
+
+        Args:
+            timestamp: Time in seconds (ignored for images)
+
+        Returns:
+            Frame as BGR numpy array (scaled to max_dimension), or None if extraction failed
+        """
+        # For images, always return the loaded image (timestamp is ignored)
+        if self._is_image:
+            return self._image_frame
+
+        # High-res video: use FFmpeg with scale filter (decodes at target res)
+        if self._use_ffmpeg_decode:
+            return self._get_frame_ffmpeg_scaled(timestamp)
+
+        if self._use_ffmpeg_fallback:
+            frame = self._get_frame_ffmpeg(timestamp)
+            return self._scale_frame(frame) if frame is not None else None
+
+        if self.cap is None:
+            return None
+
+        # Seek to timestamp
+        self.cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
+
+        ret, frame = self.cap.read()
+        if not ret:
+            # Try ffmpeg fallback for this frame
+            frame = self._get_frame_ffmpeg(timestamp)
+            return self._scale_frame(frame) if frame is not None else None
+
+        return self._scale_frame(frame)
+
+    def get_frames_at(self, timestamps: list[float]) -> list[tuple[float, np.ndarray | None]]:
+        """Extract multiple frames at given timestamps.
+
+        More efficient than calling get_frame_at repeatedly as it
+        processes timestamps in order to minimize seeking.
+
+        Args:
+            timestamps: List of times in seconds
+
+        Returns:
+            List of (timestamp, frame) tuples
+        """
+        # Sort timestamps for efficient sequential access
+        sorted_ts = sorted(set(timestamps))
+        results: dict[float, np.ndarray | None] = {}
+
+        for ts in sorted_ts:
+            results[ts] = self.get_frame_at(ts)
+
+        # Return in original order
+        return [(ts, results.get(ts)) for ts in timestamps]
+
+    def _get_frame_ffmpeg(self, timestamp: float) -> np.ndarray | None:
+        """Extract frame using ffmpeg (fallback, no scaling)."""
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+            tmp_path = tmp.name
+
+        try:
+            cmd = [
+                "ffmpeg",
+                "-y",
+                "-ss",
+                str(timestamp),
+                "-i",
+                self.video_path,
+                "-frames:v",
+                "1",
+                "-update",
+                "1",  # Required for ffmpeg 8.x single-image output
+                "-q:v",
+                "2",
+                tmp_path,
+            ]
+            subprocess.run(cmd, capture_output=True, check=True)
+
+            if os.path.exists(tmp_path):
+                frame = cv2.imread(tmp_path)
+                return frame
+        except subprocess.CalledProcessError:
+            pass
+        finally:
+            if os.path.exists(tmp_path):
+                os.unlink(tmp_path)
+
+        return None
+
+    def _get_frame_ffmpeg_scaled(self, timestamp: float) -> np.ndarray | None:
+        """Extract frame using FFmpeg with scale filter (for high-res videos).
+
+        This is faster than decoding at full resolution and then scaling with cv2.
+        """
+        if self.max_dimension is None:
+            return self._get_frame_ffmpeg(timestamp)
+
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+            tmp_path = tmp.name
+
+        try:
+            # Scale filter that maintains aspect ratio
+            # scale=W:H:force_original_aspect_ratio=decrease
+            scale_filter = f"scale={self.max_dimension}:{self.max_dimension}" f":force_original_aspect_ratio=decrease"
+
+            cmd = [
+                "ffmpeg",
+                "-y",
+                "-ss",
+                str(timestamp),
+                "-i",
+                self.video_path,
+                "-vf",
+                scale_filter,
+                "-frames:v",
+                "1",
+                "-update",
+                "1",
+                "-q:v",
+                "2",
+                tmp_path,
+            ]
+            subprocess.run(cmd, capture_output=True, check=True)
+
+            if os.path.exists(tmp_path):
+                frame = cv2.imread(tmp_path)
+                return frame
+        except subprocess.CalledProcessError:
+            pass
+        finally:
+            if os.path.exists(tmp_path):
+                os.unlink(tmp_path)
+
+        return None
+
+    def save_frame(self, frame: np.ndarray, output_path: str, quality: int = 95) -> bool:
+        """Save frame to file.
+
+        Args:
+            frame: BGR numpy array
+            output_path: Output file path
+            quality: JPEG quality (0-100)
+
+        Returns:
+            True if saved successfully
+        """
+        try:
+            cv2.imwrite(output_path, frame, [cv2.IMWRITE_JPEG_QUALITY, quality])
+            return os.path.exists(output_path)
+        except Exception as e:
+            logger.warning(f"Failed to save frame to {output_path}: {e}")
+            return False
+
+
+def extract_frames_batch(
+    video_path: str,
+    timestamps: list[float],
+    output_dir: str | None = None,
+) -> list[tuple[float, np.ndarray | None]]:
+    """Extract multiple frames from a video.
+
+    Convenience function that handles the context manager.
+
+    Args:
+        video_path: Path to video file
+        timestamps: List of timestamps in seconds
+        output_dir: Optional directory to save frames as JPEG files
+
+    Returns:
+        List of (timestamp, frame) tuples
+    """
+    with FrameExtractor(video_path) as extractor:
+        results = extractor.get_frames_at(timestamps)
+
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+            for ts, frame in results:
+                if frame is not None:
+                    output_path = os.path.join(output_dir, f"frame_{ts:.3f}.jpg")
+                    extractor.save_frame(frame, output_path)
+
+        return results
+
+
+def get_video_duration(file_path: str) -> float:
+    """Get video duration in seconds.
+
+    Args:
+        file_path: Path to video or image file
+
+    Returns:
+        Duration in seconds, or 0 for images/unknown files
+    """
+    # Check if this is an image - images have 0 duration
+    media_type = get_media_type(file_path)
+    if media_type == MediaType.IMAGE:
+        return 0.0
+
+    # Try OpenCV first (faster)
+    cap = cv2.VideoCapture(file_path)
+    if cap.isOpened():
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
+        cap.release()
+
+        if fps > 0 and frame_count > 0:
+            return frame_count / fps
+
+    # Fallback to ffprobe
+    try:
+        cmd = [
+            "ffprobe",
+            "-v",
+            "quiet",
+            "-show_entries",
+            "format=duration",
+            "-of",
+            "default=noprint_wrappers=1:nokey=1",
+            file_path,
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        return float(result.stdout.strip())
+    except (subprocess.CalledProcessError, ValueError):
+        return 0.0

media_engine/extractors/metadata/__init__.py

@@ -0,0 +1,127 @@
+"""Modular metadata extraction.
+
+This package provides manufacturer-specific metadata extractors.
+Each manufacturer module registers itself on import.
+
+Usage:
+    from media_engine.extractors.metadata import extract_metadata
+
+    metadata = extract_metadata("/path/to/video.mp4")
+
+To add a new manufacturer:
+1. Create a new module (e.g., panasonic.py)
+2. Implement a class with detect() and extract() methods
+3. Register it using: register_extractor("panasonic", PanasonicExtractor())
+4. Import the module below to trigger registration
+
+The order of imports determines detection priority.
+Specific manufacturers should be imported before generic.
+"""
+
+import logging
+from pathlib import Path
+
+from media_engine.schemas import Metadata
+
+# Import manufacturer modules to trigger registration
+# Order matters: more specific extractors first
+from . import (
+    apple,  # noqa: F401
+    arri,  # noqa: F401
+    blackmagic,  # noqa: F401
+    camera_360,  # noqa: F401 - Insta360, QooCam, GoPro MAX, etc.
+    canon,  # noqa: F401
+    dji,  # noqa: F401
+    dv,  # noqa: F401 - DV/HDV tape formats
+    ffmpeg,  # noqa: F401
+    gopro,  # noqa: F401
+    red,  # noqa: F401
+    sony,  # noqa: F401
+    tesla,  # noqa: F401
+)
+from .base import (
+    FFPROBE_WORKERS,
+    build_base_metadata,
+    extract_keyframes,
+    get_duration_fast,
+    run_ffprobe,
+    run_ffprobe_batch,
+    shutdown_ffprobe_pool,
+)
+
+# Import and register generic fallback LAST
+from .generic import GenericExtractor
+from .registry import get_extractor, list_extractors, register_extractor
+
+register_extractor("generic", GenericExtractor())
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "extract_metadata",
+    "get_duration_fast",
+    "run_ffprobe_batch",
+    "list_extractors",
+    "FFPROBE_WORKERS",
+    "shutdown_ffprobe_pool",
+]
+
+
+def extract_metadata(file_path: str, probe_data: dict | None = None) -> Metadata:
+    """Extract metadata from video file.
+
+    This function:
+    1. Runs ffprobe to get basic metadata (or uses provided probe_data)
+    2. Detects the manufacturer/device
+    3. Calls the appropriate extractor for enhanced metadata
+
+    Args:
+        file_path: Path to video file
+        probe_data: Optional pre-fetched ffprobe data (for batch processing)
+
+    Returns:
+        Metadata object with video information
+    """
+    path = Path(file_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Video file not found: {file_path}")
+
+    # Handle files that ffprobe cannot read (e.g., RED R3D)
+    # These formats require direct header parsing
+    ffprobe_unsupported = path.suffix.upper() in (".R3D",)
+
+    # Run ffprobe if not provided (and file format is supported)
+    if probe_data is None:
+        if ffprobe_unsupported:
+            # Create minimal probe_data for formats ffprobe can't read
+            probe_data = {"streams": [], "format": {"filename": file_path}}
+            logger.info(f"Skipping ffprobe for unsupported format: {path.suffix}")
+        else:
+            probe_data = run_ffprobe(file_path)
+
+    # Build base metadata (device-agnostic)
+    base_metadata = build_base_metadata(probe_data, file_path)
+
+    # Find and run the appropriate extractor
+    match = get_extractor(probe_data, file_path)
+
+    if match:
+        name, extractor = match
+        logger.info(f"Using {name} extractor for {path.name}")
+        try:
+            result = extractor.extract(probe_data, file_path, base_metadata)
+        except Exception as e:
+            logger.warning(f"Extractor {name} failed: {e}, using base metadata")
+            result = base_metadata
+    else:
+        # This shouldn't happen since generic always matches
+        logger.warning(f"No extractor matched for {path.name}")
+        result = base_metadata
+
+    # Extract keyframes (separate ffprobe call, fast with -skip_frame nokey)
+    # Done after extractor so it's not lost when extractor returns new Metadata
+    keyframes = extract_keyframes(file_path)
+    if keyframes:
+        result.keyframes = keyframes
+
+    return result

media_engine/extractors/metadata/apple.py

@@ -0,0 +1,169 @@
+"""Apple metadata extraction.
+
+Handles Apple devices:
+- iPhone (all models)
+- iPad (all models)
+- Mac (FaceTime camera, etc.)
+
+Detection methods:
+- make tag: "Apple"
+- com.apple.quicktime.make tag
+- Model contains "iPhone" or "iPad"
+
+Apple QuickTime metadata tags:
+- com.apple.quicktime.make
+- com.apple.quicktime.model
+- com.apple.quicktime.software
+- com.apple.quicktime.creationdate
+- com.apple.quicktime.location.iso6709
+"""
+
+import logging
+import re
+from typing import Any
+
+from media_engine.schemas import (
+    GPS,
+    DetectionMethod,
+    DeviceInfo,
+    MediaDeviceType,
+    Metadata,
+)
+
+from .registry import get_tags_lower, register_extractor
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_apple_location(location: str) -> GPS | None:
+    """Parse Apple's ISO 6709 location string.
+
+    Format: +59.7441+010.2045+125.0/
+    """
+    pattern = r"([+-]\d+\.?\d*)"
+    matches = re.findall(pattern, location)
+
+    if len(matches) >= 2:
+        try:
+            return GPS(
+                latitude=float(matches[0]),
+                longitude=float(matches[1]),
+                altitude=float(matches[2]) if len(matches) >= 3 else None,
+            )
+        except ValueError:
+            pass
+
+    return None
+
+
+def _determine_device_type(model: str | None) -> MediaDeviceType:
+    """Determine Apple device type from model string."""
+    if not model:
+        return MediaDeviceType.PHONE
+
+    model_upper = model.upper()
+
+    if "IPHONE" in model_upper:
+        return MediaDeviceType.PHONE
+    elif "IPAD" in model_upper:
+        return MediaDeviceType.PHONE  # Tablets are close to phones
+    elif "MAC" in model_upper or "IMAC" in model_upper or "MACBOOK" in model_upper:
+        return MediaDeviceType.CAMERA  # Mac webcams are cameras
+    else:
+        return MediaDeviceType.PHONE
+
+
+def _clean_model_name(model: str | None) -> str | None:
+    """Clean up Apple model name for display.
+
+    Examples:
+        "iPhone 15 Pro Max" -> "iPhone 15 Pro Max"
+        "iPhone15,3" -> "iPhone 15 Pro Max" (if we had a lookup table)
+    """
+    if not model:
+        return None
+
+    # Apple sometimes uses internal model identifiers like "iPhone15,3"
+    # For now, just return as-is. A future enhancement could add a lookup table.
+    return model
+
+
+class AppleExtractor:
+    """Metadata extractor for Apple devices."""
+
+    def detect(self, probe_data: dict[str, Any], file_path: str) -> bool:
+        """Detect if file is from an Apple device."""
+        tags = get_tags_lower(probe_data)
+
+        # Check make tag
+        make = tags.get("make") or tags.get("com.apple.quicktime.make")
+        if make and "APPLE" in make.upper():
+            return True
+
+        # Check model for iPhone/iPad
+        model = tags.get("model") or tags.get("com.apple.quicktime.model")
+        if model:
+            model_upper = model.upper()
+            if "IPHONE" in model_upper or "IPAD" in model_upper:
+                return True
+
+        # Check for Apple QuickTime-specific tags
+        if tags.get("com.apple.quicktime.creationdate"):
+            # This is a strong indicator of Apple origin
+            if tags.get("com.apple.quicktime.make") or tags.get("com.apple.quicktime.model"):
+                return True
+
+        return False
+
+    def extract(self, probe_data: dict[str, Any], file_path: str, base_metadata: Metadata) -> Metadata:
+        """Extract Apple-specific metadata."""
+        tags = get_tags_lower(probe_data)
+
+        # Get device info from QuickTime tags (preferred) or standard tags
+        make = tags.get("com.apple.quicktime.make") or tags.get("make") or "Apple"
+        model = tags.get("com.apple.quicktime.model") or tags.get("model")
+        software = tags.get("com.apple.quicktime.software") or tags.get("software")
+
+        # Clean up model name
+        model = _clean_model_name(model)
+
+        # Determine device type
+        device_type = _determine_device_type(model)
+
+        device = DeviceInfo(
+            make=make,
+            model=model,
+            software=software,
+            type=device_type,
+            detection_method=DetectionMethod.METADATA,
+            confidence=1.0,
+        )
+
+        # Extract GPS from Apple-specific location tag
+        gps = base_metadata.gps
+        apple_location = tags.get("com.apple.quicktime.location.iso6709")
+        if apple_location:
+            parsed_gps = _parse_apple_location(apple_location)
+            if parsed_gps:
+                gps = parsed_gps
+
+        return Metadata(
+            duration=base_metadata.duration,
+            resolution=base_metadata.resolution,
+            codec=base_metadata.codec,
+            video_codec=base_metadata.video_codec,
+            audio=base_metadata.audio,
+            fps=base_metadata.fps,
+            bitrate=base_metadata.bitrate,
+            file_size=base_metadata.file_size,
+            timecode=base_metadata.timecode,
+            created_at=base_metadata.created_at,
+            device=device,
+            gps=gps,
+            color_space=base_metadata.color_space,
+            lens=base_metadata.lens,
+        )
+
+
+# Register this extractor
+register_extractor("apple", AppleExtractor())