Extract image dimensions into asset metadata (#13991)

Amp-Thread-ID: https://ampcode.com/threads/T-019e4ca5-b71a-7168-8f56-58b2325f34c3
Co-authored-by: Amp <amp@ampcode.com>

app/assets/scanner.py

@@ -33,6 +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.metadata_extract import extract_file_metadata
 from app.assets.services.path_utils import (
     compute_relative_filename,
@@ -506,6 +507,10 @@ 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)
         set_reference_system_metadata(session, reference_id, system_metadata)
 
     if full_hash:

app/assets/services/image_dimensions.py

@@ -0,0 +1,63 @@
+"""Image dimension extraction for asset ingest.
+
+Reads only the image header via Pillow to capture width/height cheaply,
+without a full pixel decode. 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_image_dimensions(
+    file_path: str, mime_type: str | None = None
+) -> dict[str, Any] | None:
+    """Extract image dimensions 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 ``image/``, extraction is skipped without touching the file.
+
+    Returns:
+        ``{"kind": "image", "width": W, "height": H}`` when the file is a
+        recognizable image with positive dimensions, otherwise ``None``.
+
+    The dict shape is intended to be merged into ``system_metadata`` so the
+    asset response surfaces ``metadata.kind`` plus dimension fields for image
+    assets. Forward-compatible: future media kinds (e.g. ``"video"`` with
+    duration/fps) can extend this shape without schema changes.
+    """
+    if mime_type is not None and not mime_type.startswith("image/"):
+        return None
+
+    try:
+        from PIL import Image, UnidentifiedImageError
+    except ImportError:
+        logger.debug(
+            "Pillow not available; skipping image dimension extraction for %s",
+            file_path,
+        )
+        return None
+
+    try:
+        with Image.open(file_path) as img:
+            width, height = img.size
+    except (OSError, UnidentifiedImageError, ValueError) as exc:
+        logger.debug(
+            "Failed to read image dimensions 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
+
+    return {"kind": "image", "width": width, "height": height}

app/assets/services/ingest.py

@@ -17,9 +17,11 @@ from app.assets.database.queries import (
     get_reference_by_file_path,
     get_reference_tags,
     get_or_create_reference,
+    list_references_by_asset_id,
     reference_exists,
     remove_missing_tag_for_asset_id,
     set_reference_metadata,
+    set_reference_system_metadata,
     set_reference_tags,
     update_asset_hash_and_mime,
     upsert_asset,
@@ -29,6 +31,7 @@ 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.path_utils import (
     compute_relative_filename,
     get_name_and_tags_from_asset_path,
@@ -118,6 +121,14 @@ def _ingest_file_from_path(
                 user_metadata=user_metadata,
             )
 
+            _maybe_store_image_dimensions(
+                session,
+                reference_id=reference_id,
+                file_path=locator,
+                mime_type=mime_type,
+                current_system_metadata=ref.system_metadata,
+            )
+
         try:
             remove_missing_tag_for_asset_id(session, asset_id=asset.id)
         except Exception:
@@ -288,6 +299,13 @@ def _register_existing_asset(
                 user_metadata=new_meta,
             )
 
+        _backfill_image_dimensions_from_siblings(
+            session,
+            asset_id=asset.id,
+            new_reference_id=ref.id,
+            current_system_metadata=ref.system_metadata,
+        )
+
         if tags is not None:
             set_reference_tags(
                 session,
@@ -334,6 +352,87 @@ def _update_metadata_with_filename(
         )
 
 
+_IMAGE_DIMENSION_KEYS = ("kind", "width", "height")
+
+
+def _maybe_store_image_dimensions(
+    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.
+
+    Non-image 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)
+    if not dims:
+        return
+
+    current = current_system_metadata or {}
+    merged = dict(current)
+    merged.update(dims)
+    if merged != current:
+        set_reference_system_metadata(
+            session,
+            reference_id=reference_id,
+            system_metadata=merged,
+        )
+
+
+def _backfill_image_dimensions_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.
+
+    The from-hash path doesn't read the file bytes, so dimensions 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.
+
+    Best-effort: missing siblings, non-image siblings, or absent dimension
+    keys leave the target reference unchanged.
+    """
+    current = current_system_metadata or {}
+    if current.get("kind") == "image" 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":
+            continue
+        width = meta.get("width")
+        height = meta.get("height")
+        if (
+            type(width) is not int
+            or type(height) is not int
+            or width <= 0
+            or height <= 0
+        ):
+            continue
+        merged = dict(current)
+        merged["kind"] = "image"
+        merged["width"] = width
+        merged["height"] = height
+        if merged != current:
+            set_reference_system_metadata(
+                session,
+                reference_id=new_reference_id,
+                system_metadata=merged,
+            )
+        return
+
+
 def _sanitize_filename(name: str | None, fallback: str) -> str:
     n = os.path.basename((name or "").strip() or fallback)
     return n if n else fallback