diff --git a/app/assets/scanner.py b/app/assets/scanner.py index 42c4c1e9d..17a445d56 100644 --- a/app/assets/scanner.py +++ b/app/assets/scanner.py @@ -33,7 +33,7 @@ from app.assets.services.file_utils import ( verify_file_unchanged, ) from app.assets.services.hashing import HashCheckpoint, compute_blake3_hash -from app.assets.services.image_dimensions import extract_image_dimensions +from app.assets.services.media_metadata import extract_media_metadata from app.assets.services.metadata_extract import extract_file_metadata from app.assets.services.path_utils import ( compute_loader_path, @@ -507,10 +507,9 @@ def enrich_asset( if extract_metadata and metadata: system_metadata = metadata.to_user_metadata() - if mime_type and mime_type.startswith("image/"): - dims = extract_image_dimensions(file_path, mime_type=mime_type) - if dims: - system_metadata.update(dims) + dims = extract_media_metadata(file_path, mime_type=mime_type) + if dims: + system_metadata.update(dims) set_reference_system_metadata(session, reference_id, system_metadata) if full_hash: diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index 1ffb3d634..bbdfe9f1e 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -31,7 +31,10 @@ from app.assets.database.queries import ( from app.assets.helpers import get_utc_now, normalize_tags from app.assets.services.bulk_ingest import batch_insert_seed_assets from app.assets.services.file_utils import get_size_and_mtime_ns -from app.assets.services.image_dimensions import extract_image_dimensions +from app.assets.services.media_metadata import ( + MEDIA_METADATA_KEYS, + extract_media_metadata, +) from app.assets.services.path_utils import ( compute_loader_path, get_name_and_tags_from_asset_path, @@ -138,7 +141,7 @@ def _ingest_file_from_path( user_metadata=user_metadata, ) - _maybe_store_image_dimensions( + _maybe_store_media_metadata( session, reference_id=reference_id, file_path=locator, @@ -316,7 +319,7 @@ def _register_existing_asset( user_metadata=new_meta, ) - _backfill_image_dimensions_from_siblings( + _backfill_media_metadata_from_siblings( session, asset_id=asset.id, new_reference_id=ref.id, @@ -369,25 +372,22 @@ def _update_metadata_with_filename( ) -_IMAGE_DIMENSION_KEYS = ("kind", "width", "height") +_MEDIA_KINDS = ("image", "video") -def _maybe_store_image_dimensions( +def _maybe_store_media_metadata( session: Session, reference_id: str, file_path: str, mime_type: str | None, current_system_metadata: dict | None, ) -> None: - """Populate ``kind``/``width``/``height`` on system_metadata for image refs. + """Populate media keys on system_metadata for image and video refs. - Non-image MIME types are a no-op. Pre-existing keys (e.g. enricher-written + Non-media MIME types are a no-op. Pre-existing keys (e.g. enricher-written safetensors metadata, download provenance) are preserved by merge. """ - if not mime_type or not mime_type.startswith("image/"): - return - - dims = extract_image_dimensions(file_path, mime_type=mime_type) + dims = extract_media_metadata(file_path, mime_type=mime_type) if not dims: return @@ -402,31 +402,35 @@ def _maybe_store_image_dimensions( ) -def _backfill_image_dimensions_from_siblings( +def _backfill_media_metadata_from_siblings( session: Session, asset_id: str, new_reference_id: str, current_system_metadata: dict | None, ) -> None: - """Copy image dimension keys from any sibling reference of the same asset. + """Copy media metadata keys from any sibling reference of the same asset. - The from-hash path doesn't read the file bytes, so dimensions can't be + The from-hash path doesn't read the file bytes, so metadata can't be extracted there directly. When another reference of the same asset already - carries image dimensions, copy them onto the new reference so consumers - see consistent metadata regardless of how the asset was registered. + carries media metadata, copy it onto the new reference so consumers see + consistent metadata regardless of how the asset was registered. - Best-effort: missing siblings, non-image siblings, or absent dimension + Best-effort: missing siblings, non-media siblings, or absent metadata keys leave the target reference unchanged. """ current = current_system_metadata or {} - if current.get("kind") == "image" and "width" in current and "height" in current: + if ( + current.get("kind") in _MEDIA_KINDS + and "width" in current + and "height" in current + ): return for sibling in list_references_by_asset_id(session, asset_id): if sibling.id == new_reference_id: continue meta = sibling.system_metadata or {} - if meta.get("kind") != "image": + if meta.get("kind") not in _MEDIA_KINDS: continue width = meta.get("width") height = meta.get("height") @@ -438,9 +442,9 @@ def _backfill_image_dimensions_from_siblings( ): continue merged = dict(current) - merged["kind"] = "image" - merged["width"] = width - merged["height"] = height + for key in MEDIA_METADATA_KEYS: + if key in meta: + merged[key] = meta[key] if merged != current: set_reference_system_metadata( session, diff --git a/app/assets/services/media_metadata.py b/app/assets/services/media_metadata.py new file mode 100644 index 000000000..ca0ddd539 --- /dev/null +++ b/app/assets/services/media_metadata.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from typing import Any + +from app.assets.services.image_dimensions import extract_image_dimensions +from app.assets.services.video_metadata import extract_video_metadata + +MEDIA_METADATA_KEYS = ( + "kind", + "width", + "height", + "duration", + "fps", + "frame_count", +) + + +def extract_media_metadata( + file_path: str, mime_type: str | None = None +) -> dict[str, Any] | None: + """Extract media metadata for the file, dispatched on the MIME prefix. + + Returns ``None`` for non-media MIME types or unreadable files. + """ + if not mime_type: + return None + if mime_type.startswith("image/"): + return extract_image_dimensions(file_path, mime_type=mime_type) + if mime_type.startswith("video/"): + return extract_video_metadata(file_path, mime_type=mime_type) + return None diff --git a/app/assets/services/video_metadata.py b/app/assets/services/video_metadata.py new file mode 100644 index 000000000..603b503c0 --- /dev/null +++ b/app/assets/services/video_metadata.py @@ -0,0 +1,81 @@ +"""Video metadata extraction for asset ingest. + +Reads only the container/stream headers via PyAV to capture dimensions, +duration, fps and frame count cheaply, without decoding frames. Returns a +metadata dict suitable for merging into ``AssetReference.system_metadata``. +""" +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def extract_video_metadata( + file_path: str, mime_type: str | None = None +) -> dict[str, Any] | None: + """Extract video stream metadata for the file at ``file_path``. + + Args: + file_path: Absolute path to a file on disk. + mime_type: Optional MIME type hint. When provided and not prefixed + with ``video/``, extraction is skipped without touching the file. + + Returns: + ``{"kind": "video", "width": W, "height": H, ...}`` with optional + ``duration`` (seconds), ``fps`` and ``frame_count`` keys when the + file has a recognizable video stream, otherwise ``None``. + """ + if mime_type is not None and not mime_type.startswith("video/"): + return None + + try: + import av + except ImportError: + logger.debug( + "PyAV not available; skipping video metadata extraction for %s", + file_path, + ) + return None + + try: + with av.open(file_path) as container: + stream = next( + (s for s in container.streams if s.type == "video"), None + ) + if stream is None: + return None + + fps = float(stream.average_rate) if stream.average_rate else None + duration = None + if stream.duration is not None and stream.time_base is not None: + duration = float(stream.duration * stream.time_base) + elif container.duration is not None: + duration = float(container.duration * av.time_base) + frame_count = stream.frames or None + if frame_count is None and duration is not None and fps is not None: + frame_count = round(duration * fps) + + width = stream.codec_context.width + height = stream.codec_context.height + except (OSError, ValueError, av.error.FFmpegError) as exc: + logger.debug("Failed to read video metadata from %s: %s", file_path, exc) + return None + + if ( + not isinstance(width, int) + or not isinstance(height, int) + or width <= 0 + or height <= 0 + ): + return None + + metadata: dict[str, Any] = {"kind": "video", "width": width, "height": height} + if duration is not None: + metadata["duration"] = duration + if fps is not None: + metadata["fps"] = fps + if frame_count is not None: + metadata["frame_count"] = frame_count + return metadata diff --git a/tests-unit/assets_test/services/test_video_metadata.py b/tests-unit/assets_test/services/test_video_metadata.py new file mode 100644 index 000000000..cb9be0038 --- /dev/null +++ b/tests-unit/assets_test/services/test_video_metadata.py @@ -0,0 +1,138 @@ +"""Tests for the video_metadata service.""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from app.assets.services.media_metadata import extract_media_metadata +from app.assets.services.video_metadata import extract_video_metadata + +av = pytest.importorskip("av") + + +def _make_mp4( + path: Path, width: int = 64, height: int = 48, frames: int = 12, fps: int = 8 +) -> Path: + with av.open(str(path), "w") as container: + stream = container.add_stream("libx264", rate=fps) + stream.width = width + stream.height = height + stream.pix_fmt = "yuv420p" + for _ in range(frames): + frame = av.VideoFrame.from_ndarray( + np.zeros((height, width, 3), dtype=np.uint8), format="rgb24" + ) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return path + + +class TestExtractVideoMetadata: + def test_extracts_stream_metadata(self, tmp_path: Path): + f = _make_mp4(tmp_path / "clip.mp4", width=64, height=48, frames=12, fps=8) + + result = extract_video_metadata(str(f), mime_type="video/mp4") + + assert result is not None + assert result["kind"] == "video" + assert result["width"] == 64 + assert result["height"] == 48 + assert result["frame_count"] == 12 + assert result["fps"] == pytest.approx(8.0) + assert result["duration"] == pytest.approx(12 / 8, abs=0.1) + + def test_works_when_mime_type_is_none(self, tmp_path: Path): + f = _make_mp4(tmp_path / "no_mime.mp4") + + result = extract_video_metadata(str(f), mime_type=None) + + assert result is not None + assert result["kind"] == "video" + + @pytest.mark.parametrize( + "mime", + ["application/json", "text/plain", "image/png", "audio/mpeg"], + ) + def test_skips_non_video_mime_types(self, tmp_path: Path, mime: str): + result = extract_video_metadata( + str(tmp_path / "untouched.mp4"), mime_type=mime + ) + + assert result is None + + def test_returns_none_for_missing_file(self, tmp_path: Path): + result = extract_video_metadata( + str(tmp_path / "does_not_exist.mp4"), mime_type="video/mp4" + ) + + assert result is None + + def test_returns_none_for_corrupt_video(self, tmp_path: Path): + f = tmp_path / "corrupt.mp4" + f.write_bytes(b"not actually an mp4 file") + + result = extract_video_metadata(str(f), mime_type="video/mp4") + + assert result is None + + +class TestExtractMediaMetadata: + def test_dispatches_video_mime_to_video_extractor(self, tmp_path: Path): + f = _make_mp4(tmp_path / "clip.mp4") + + result = extract_media_metadata(str(f), mime_type="video/mp4") + + assert result is not None + assert result["kind"] == "video" + + def test_dispatches_image_mime_to_image_extractor(self, tmp_path: Path): + from PIL import Image + + f = tmp_path / "img.png" + Image.new("RGB", (32, 16)).save(f, format="PNG") + + result = extract_media_metadata(str(f), mime_type="image/png") + + assert result == {"kind": "image", "width": 32, "height": 16} + + def test_returns_none_without_mime_type(self, tmp_path: Path): + f = _make_mp4(tmp_path / "clip.mp4") + + assert extract_media_metadata(str(f), mime_type=None) is None + + def test_returns_none_for_non_media_mime(self, tmp_path: Path): + f = tmp_path / "file.bin" + f.write_bytes(b"\x00") + + assert extract_media_metadata(str(f), mime_type="text/plain") is None + + +class TestIngestStoresVideoMetadata: + def test_register_file_in_place_stores_video_metadata( + self, mock_create_session, temp_dir: Path, session + ): + from app.assets.database.models import AssetReference + from app.assets.services.ingest import _ingest_file_from_path + + f = _make_mp4(temp_dir / "clip.mp4", width=64, height=48) + + result = _ingest_file_from_path( + abs_path=str(f), + asset_hash="blake3:video123", + size_bytes=f.stat().st_size, + mtime_ns=1234567890000000000, + mime_type="video/mp4", + ) + + assert result.reference_id is not None + ref = session.query(AssetReference).one() + meta = ref.system_metadata or {} + assert meta["kind"] == "video" + assert meta["width"] == 64 + assert meta["height"] == 48 + assert meta["frame_count"] == 12 + assert meta["fps"] == pytest.approx(8.0)