mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-25 18:32:35 +08:00
feat(assets): extract video metadata into system_metadata on ingest and scan
This commit is contained in:
@@ -33,7 +33,7 @@ from app.assets.services.file_utils import (
|
|||||||
verify_file_unchanged,
|
verify_file_unchanged,
|
||||||
)
|
)
|
||||||
from app.assets.services.hashing import HashCheckpoint, compute_blake3_hash
|
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.metadata_extract import extract_file_metadata
|
||||||
from app.assets.services.path_utils import (
|
from app.assets.services.path_utils import (
|
||||||
compute_loader_path,
|
compute_loader_path,
|
||||||
@@ -507,10 +507,9 @@ def enrich_asset(
|
|||||||
|
|
||||||
if extract_metadata and metadata:
|
if extract_metadata and metadata:
|
||||||
system_metadata = metadata.to_user_metadata()
|
system_metadata = metadata.to_user_metadata()
|
||||||
if mime_type and mime_type.startswith("image/"):
|
dims = extract_media_metadata(file_path, mime_type=mime_type)
|
||||||
dims = extract_image_dimensions(file_path, mime_type=mime_type)
|
if dims:
|
||||||
if dims:
|
system_metadata.update(dims)
|
||||||
system_metadata.update(dims)
|
|
||||||
set_reference_system_metadata(session, reference_id, system_metadata)
|
set_reference_system_metadata(session, reference_id, system_metadata)
|
||||||
|
|
||||||
if full_hash:
|
if full_hash:
|
||||||
|
|||||||
@@ -31,7 +31,10 @@ from app.assets.database.queries import (
|
|||||||
from app.assets.helpers import get_utc_now, normalize_tags
|
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.bulk_ingest import batch_insert_seed_assets
|
||||||
from app.assets.services.file_utils import get_size_and_mtime_ns
|
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 (
|
from app.assets.services.path_utils import (
|
||||||
compute_loader_path,
|
compute_loader_path,
|
||||||
get_name_and_tags_from_asset_path,
|
get_name_and_tags_from_asset_path,
|
||||||
@@ -138,7 +141,7 @@ def _ingest_file_from_path(
|
|||||||
user_metadata=user_metadata,
|
user_metadata=user_metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
_maybe_store_image_dimensions(
|
_maybe_store_media_metadata(
|
||||||
session,
|
session,
|
||||||
reference_id=reference_id,
|
reference_id=reference_id,
|
||||||
file_path=locator,
|
file_path=locator,
|
||||||
@@ -316,7 +319,7 @@ def _register_existing_asset(
|
|||||||
user_metadata=new_meta,
|
user_metadata=new_meta,
|
||||||
)
|
)
|
||||||
|
|
||||||
_backfill_image_dimensions_from_siblings(
|
_backfill_media_metadata_from_siblings(
|
||||||
session,
|
session,
|
||||||
asset_id=asset.id,
|
asset_id=asset.id,
|
||||||
new_reference_id=ref.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,
|
session: Session,
|
||||||
reference_id: str,
|
reference_id: str,
|
||||||
file_path: str,
|
file_path: str,
|
||||||
mime_type: str | None,
|
mime_type: str | None,
|
||||||
current_system_metadata: dict | None,
|
current_system_metadata: dict | None,
|
||||||
) -> 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.
|
safetensors metadata, download provenance) are preserved by merge.
|
||||||
"""
|
"""
|
||||||
if not mime_type or not mime_type.startswith("image/"):
|
dims = extract_media_metadata(file_path, mime_type=mime_type)
|
||||||
return
|
|
||||||
|
|
||||||
dims = extract_image_dimensions(file_path, mime_type=mime_type)
|
|
||||||
if not dims:
|
if not dims:
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -402,31 +402,35 @@ def _maybe_store_image_dimensions(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _backfill_image_dimensions_from_siblings(
|
def _backfill_media_metadata_from_siblings(
|
||||||
session: Session,
|
session: Session,
|
||||||
asset_id: str,
|
asset_id: str,
|
||||||
new_reference_id: str,
|
new_reference_id: str,
|
||||||
current_system_metadata: dict | None,
|
current_system_metadata: dict | None,
|
||||||
) -> 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
|
extracted there directly. When another reference of the same asset already
|
||||||
carries image dimensions, copy them onto the new reference so consumers
|
carries media metadata, copy it onto the new reference so consumers see
|
||||||
see consistent metadata regardless of how the asset was registered.
|
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.
|
keys leave the target reference unchanged.
|
||||||
"""
|
"""
|
||||||
current = current_system_metadata or {}
|
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
|
return
|
||||||
|
|
||||||
for sibling in list_references_by_asset_id(session, asset_id):
|
for sibling in list_references_by_asset_id(session, asset_id):
|
||||||
if sibling.id == new_reference_id:
|
if sibling.id == new_reference_id:
|
||||||
continue
|
continue
|
||||||
meta = sibling.system_metadata or {}
|
meta = sibling.system_metadata or {}
|
||||||
if meta.get("kind") != "image":
|
if meta.get("kind") not in _MEDIA_KINDS:
|
||||||
continue
|
continue
|
||||||
width = meta.get("width")
|
width = meta.get("width")
|
||||||
height = meta.get("height")
|
height = meta.get("height")
|
||||||
@@ -438,9 +442,9 @@ def _backfill_image_dimensions_from_siblings(
|
|||||||
):
|
):
|
||||||
continue
|
continue
|
||||||
merged = dict(current)
|
merged = dict(current)
|
||||||
merged["kind"] = "image"
|
for key in MEDIA_METADATA_KEYS:
|
||||||
merged["width"] = width
|
if key in meta:
|
||||||
merged["height"] = height
|
merged[key] = meta[key]
|
||||||
if merged != current:
|
if merged != current:
|
||||||
set_reference_system_metadata(
|
set_reference_system_metadata(
|
||||||
session,
|
session,
|
||||||
|
|||||||
31
app/assets/services/media_metadata.py
Normal file
31
app/assets/services/media_metadata.py
Normal file
@@ -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
|
||||||
81
app/assets/services/video_metadata.py
Normal file
81
app/assets/services/video_metadata.py
Normal file
@@ -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
|
||||||
138
tests-unit/assets_test/services/test_video_metadata.py
Normal file
138
tests-unit/assets_test/services/test_video_metadata.py
Normal file
@@ -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)
|
||||||
Reference in New Issue
Block a user