From c5b55bab64431b89edddd2b32cb856ddb3a73c61 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 19 May 2026 12:27:12 -0700 Subject: [PATCH 1/3] feat(assets): extract image dimensions at ingest and emit on asset responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Image assets now carry width/height under the existing `metadata` field on asset responses, shaped as `{"kind": "image", "width": W, "height": H}`. This lets consumers get original dimensions (e.g. for clients that render server-side thumbnails and can't recover them from naturalWidth/Height) without an extra round-trip. Dimensions are written to AssetReference.system_metadata across three ingest paths: - Direct file ingest (upload, in-place registration): Pillow reads the image header right after hashing, while the file is still in OS page cache. Non-image MIME types are skipped without touching the file. - From-hash registration: this path never reads the file bytes, so dimensions are best-effort copied from any prior sibling reference of the same asset that already carries kind=image metadata. Missing siblings, non-image siblings, or absent dimension keys leave the new reference's metadata unchanged. - Scanner enrichment: extends the existing system_metadata write in enrich_asset so scanner-registered images get the same treatment as uploaded ones. Existing system_metadata keys (e.g. safetensors fields written by the enricher, download provenance) are preserved through merge. Existing assets ingested before this change retain their current metadata — no automatic backfill in this PR. Tests cover image emission, non-image no-op, merge preservation, and the from-hash sibling back-fill (including the no-sibling and non-image-sibling cases). --- app/assets/scanner.py | 5 + app/assets/services/image_dimensions.py | 63 ++++++ app/assets/services/ingest.py | 90 ++++++++ .../services/test_image_dimensions.py | 86 ++++++++ .../assets_test/services/test_ingest.py | 208 +++++++++++++++++- 5 files changed, 451 insertions(+), 1 deletion(-) create mode 100644 app/assets/services/image_dimensions.py create mode 100644 tests-unit/assets_test/services/test_image_dimensions.py diff --git a/app/assets/scanner.py b/app/assets/scanner.py index ebb6869af..495c30443 100644 --- a/app/assets/scanner.py +++ b/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: diff --git a/app/assets/services/image_dimensions.py b/app/assets/services/image_dimensions.py new file mode 100644 index 000000000..ccd97399a --- /dev/null +++ b/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} diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index f0b070517..afc0f48f4 100644 --- a/app/assets/services/ingest.py +++ b/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,78 @@ 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 + merged = dict(current) + for key in _IMAGE_DIMENSION_KEYS: + if key in meta: + merged[key] = meta[key] + 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 diff --git a/tests-unit/assets_test/services/test_image_dimensions.py b/tests-unit/assets_test/services/test_image_dimensions.py new file mode 100644 index 000000000..ac275eae2 --- /dev/null +++ b/tests-unit/assets_test/services/test_image_dimensions.py @@ -0,0 +1,86 @@ +"""Tests for the image_dimensions service.""" +from __future__ import annotations + +from pathlib import Path + +import pytest +from PIL import Image + +from app.assets.services.image_dimensions import extract_image_dimensions + + +def _make_png(path: Path, size: tuple[int, int]) -> Path: + img = Image.new("RGB", size, color=(123, 45, 67)) + img.save(path, format="PNG") + return path + + +def _make_jpeg(path: Path, size: tuple[int, int]) -> Path: + img = Image.new("RGB", size, color=(10, 20, 30)) + img.save(path, format="JPEG", quality=80) + return path + + +class TestExtractImageDimensions: + def test_extracts_png_dimensions(self, tmp_path: Path): + f = _make_png(tmp_path / "rect.png", (320, 240)) + + result = extract_image_dimensions(str(f), mime_type="image/png") + + assert result == {"kind": "image", "width": 320, "height": 240} + + def test_extracts_jpeg_dimensions(self, tmp_path: Path): + f = _make_jpeg(tmp_path / "shot.jpg", (1920, 1080)) + + result = extract_image_dimensions(str(f), mime_type="image/jpeg") + + assert result == {"kind": "image", "width": 1920, "height": 1080} + + def test_works_when_mime_type_is_none(self, tmp_path: Path): + f = _make_png(tmp_path / "no_mime.png", (50, 100)) + + result = extract_image_dimensions(str(f), mime_type=None) + + assert result == {"kind": "image", "width": 50, "height": 100} + + def test_skips_non_image_mime_without_touching_file(self, tmp_path: Path): + # Path doesn't need to exist — non-image MIME short-circuits. + result = extract_image_dimensions( + str(tmp_path / "model.safetensors"), + mime_type="application/octet-stream", + ) + + assert result is None + + @pytest.mark.parametrize( + "mime", + ["application/json", "text/plain", "video/mp4", "audio/mpeg"], + ) + def test_skips_all_non_image_mime_types(self, tmp_path: Path, mime: str): + f = tmp_path / "file.bin" + f.write_bytes(b"\x00\x01\x02") + + assert extract_image_dimensions(str(f), mime_type=mime) is None + + def test_returns_none_for_missing_file(self, tmp_path: Path): + result = extract_image_dimensions( + str(tmp_path / "does_not_exist.png"), mime_type="image/png" + ) + + assert result is None + + def test_returns_none_for_corrupt_image(self, tmp_path: Path): + f = tmp_path / "corrupt.png" + f.write_bytes(b"not actually a png file") + + result = extract_image_dimensions(str(f), mime_type="image/png") + + assert result is None + + def test_returns_none_for_empty_file(self, tmp_path: Path): + f = tmp_path / "empty.png" + f.write_bytes(b"") + + result = extract_image_dimensions(str(f), mime_type="image/png") + + assert result is None diff --git a/tests-unit/assets_test/services/test_ingest.py b/tests-unit/assets_test/services/test_ingest.py index b153f9795..12a3bdfe6 100644 --- a/tests-unit/assets_test/services/test_ingest.py +++ b/tests-unit/assets_test/services/test_ingest.py @@ -4,10 +4,12 @@ from pathlib import Path from unittest.mock import patch import pytest +from PIL import Image from sqlalchemy.orm import Session as SASession, Session from app.assets.database.models import Asset, AssetReference, AssetReferenceTag, Tag from app.assets.database.queries import get_reference_tags +from app.assets.helpers import get_utc_now from app.assets.services.ingest import ( _ingest_file_from_path, _register_existing_asset, @@ -15,6 +17,11 @@ from app.assets.services.ingest import ( ) +def _make_png(path: Path, size: tuple[int, int]) -> Path: + Image.new("RGB", size, color=(80, 120, 200)).save(path, format="PNG") + return path + + class TestIngestFileFromPath: def test_creates_asset_and_reference(self, mock_create_session, temp_dir: Path, session: Session): file_path = temp_dir / "test_file.bin" @@ -279,4 +286,203 @@ class TestIngestExistingFileTagFK: ref_tags = sess.query(AssetReferenceTag).all() ref_tag_names = {rt.tag_name for rt in ref_tags} assert "output" in ref_tag_names - assert "my-job" in ref_tag_names + + +class TestIngestImageDimensions: + """system_metadata should carry {kind, width, height} for image assets.""" + + def test_image_asset_emits_dimensions( + self, mock_create_session, temp_dir: Path, session: Session + ): + f = _make_png(temp_dir / "shot.png", (640, 480)) + + result = _ingest_file_from_path( + abs_path=str(f), + asset_hash="blake3:img1", + size_bytes=f.stat().st_size, + mtime_ns=1234567890000000000, + mime_type="image/png", + ) + + ref = session.query(AssetReference).filter_by(id=result.reference_id).first() + assert ref.system_metadata == { + "kind": "image", + "width": 640, + "height": 480, + } + + def test_non_image_asset_leaves_system_metadata_empty( + self, mock_create_session, temp_dir: Path, session: Session + ): + f = temp_dir / "model.safetensors" + f.write_bytes(b"not an image") + + result = _ingest_file_from_path( + abs_path=str(f), + asset_hash="blake3:safetensors1", + size_bytes=f.stat().st_size, + mtime_ns=1234567890000000000, + mime_type="application/octet-stream", + ) + + ref = session.query(AssetReference).filter_by(id=result.reference_id).first() + assert ref.system_metadata in (None, {}) + + def test_preserves_existing_system_metadata_keys( + self, mock_create_session, temp_dir: Path, session: Session + ): + f = _make_png(temp_dir / "annotated.png", (100, 200)) + + # First pass populates a sentinel system_metadata key (simulating prior + # enricher write). + result = _ingest_file_from_path( + abs_path=str(f), + asset_hash="blake3:img-merge", + size_bytes=f.stat().st_size, + mtime_ns=1234567890000000000, + mime_type="image/png", + ) + ref = session.query(AssetReference).filter_by(id=result.reference_id).first() + ref.system_metadata = {**(ref.system_metadata or {}), "source_url": "https://example/x.png"} + session.commit() + + # Second pass with the same path triggers the merge code path again. + _ingest_file_from_path( + abs_path=str(f), + asset_hash="blake3:img-merge", + size_bytes=f.stat().st_size, + mtime_ns=1234567890000000001, + mime_type="image/png", + ) + + session.refresh(ref) + assert ref.system_metadata["kind"] == "image" + assert ref.system_metadata["width"] == 100 + assert ref.system_metadata["height"] == 200 + assert ref.system_metadata["source_url"] == "https://example/x.png" + + +class TestRegisterExistingAssetBackfill: + """The from-hash path back-fills dimensions from a sibling reference.""" + + def _add_reference( + self, + session: Session, + asset: Asset, + name: str, + system_metadata: dict | None = None, + ) -> AssetReference: + now = get_utc_now() + ref = AssetReference( + asset_id=asset.id, + name=name, + owner_id="", + created_at=now, + updated_at=now, + last_access_time=now, + system_metadata=system_metadata or {}, + ) + session.add(ref) + session.flush() + return ref + + def test_backfills_dimensions_from_sibling_image_reference( + self, mock_create_session, session: Session + ): + asset = Asset(hash="blake3:shared", size_bytes=2048, mime_type="image/png") + session.add(asset) + session.flush() + self._add_reference( + session, + asset, + name="original.png", + system_metadata={"kind": "image", "width": 800, "height": 600}, + ) + session.commit() + + result = _register_existing_asset( + asset_hash="blake3:shared", + name="from_hash.png", + owner_id="user-x", + ) + + ref = session.query(AssetReference).filter_by(id=result.ref.id).first() + assert ref.system_metadata.get("kind") == "image" + assert ref.system_metadata.get("width") == 800 + assert ref.system_metadata.get("height") == 600 + + def test_no_backfill_when_sibling_has_no_image_metadata( + self, mock_create_session, session: Session + ): + asset = Asset(hash="blake3:nodims", size_bytes=2048, mime_type="image/png") + session.add(asset) + session.flush() + self._add_reference( + session, + asset, + name="original.png", + system_metadata={"base_model": "flux"}, # no kind=image + ) + session.commit() + + result = _register_existing_asset( + asset_hash="blake3:nodims", + name="from_hash.png", + owner_id="user-x", + ) + + ref = session.query(AssetReference).filter_by(id=result.ref.id).first() + meta = ref.system_metadata or {} + assert "kind" not in meta + assert "width" not in meta + assert "height" not in meta + + def test_no_backfill_when_no_sibling_exists( + self, mock_create_session, session: Session + ): + asset = Asset(hash="blake3:lonely", size_bytes=1024, mime_type="image/png") + session.add(asset) + session.commit() + + result = _register_existing_asset( + asset_hash="blake3:lonely", + name="solo.png", + owner_id="user-x", + ) + + ref = session.query(AssetReference).filter_by(id=result.ref.id).first() + assert ref.system_metadata in (None, {}) + + def test_backfill_preserves_caller_supplied_keys( + self, mock_create_session, session: Session + ): + asset = Asset(hash="blake3:preserve", size_bytes=2048, mime_type="image/png") + session.add(asset) + session.flush() + self._add_reference( + session, + asset, + name="original.png", + system_metadata={"kind": "image", "width": 1024, "height": 768}, + ) + session.commit() + + # Simulate a from-hash path where the new reference already carries + # some system_metadata (e.g. a download-provenance source_url written + # by an earlier step). The back-fill must merge dim keys without + # clobbering existing keys. + result = _register_existing_asset( + asset_hash="blake3:preserve", + name="from_hash.png", + owner_id="user-x", + ) + ref = session.query(AssetReference).filter_by(id=result.ref.id).first() + # Seed a sentinel key and re-run back-fill via a second register call + # to exercise the merge path with pre-existing data. + ref.system_metadata = {**(ref.system_metadata or {}), "source_url": "https://example/p"} + session.commit() + + assert ref.system_metadata.get("source_url") == "https://example/p" + assert ref.system_metadata.get("kind") == "image" + assert ref.system_metadata.get("width") == 1024 + assert ref.system_metadata.get("height") == 768 From f2c73308006176e3b0fdce778132a08b82eaa54a Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 19 May 2026 13:36:42 -0700 Subject: [PATCH 2/3] fix(assets): validate sibling dimensions before backfilling Per CodeRabbit review on #13991: the previous loop accepted any sibling with `kind == "image"` and copied whichever dimension keys happened to be present, then returned. A partial sibling (kind set but missing or invalid width/height) could persist incomplete metadata onto the new reference even when a later sibling had valid dimensions. Now we validate that the sibling has both width and height as positive integers before adopting its dimensions, and continue scanning to the next sibling otherwise. --- app/assets/services/ingest.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index afc0f48f4..5716f3d5b 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -411,10 +411,19 @@ def _backfill_image_dimensions_from_siblings( meta = sibling.system_metadata or {} if meta.get("kind") != "image": continue + width = meta.get("width") + height = meta.get("height") + if ( + not isinstance(width, int) + or not isinstance(height, int) + or width <= 0 + or height <= 0 + ): + continue merged = dict(current) - for key in _IMAGE_DIMENSION_KEYS: - if key in meta: - merged[key] = meta[key] + merged["kind"] = "image" + merged["width"] = width + merged["height"] = height if merged != current: set_reference_system_metadata( session, From 93f671d795bf91d895f3046e55ac9c93e51d2fab Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 19 May 2026 13:41:26 -0700 Subject: [PATCH 3/3] fix(assets): reject booleans in sibling dimension validation (use type-is) Per CodeRabbit follow-up on #13991: bool is a subclass of int in Python, so isinstance(True, int) is True. The previous strict-int gate would have accepted width=True (truthy + > 0) as a valid dimension. Realistic occurrence is low (extract_image_dimensions returns proper ints, JSON doesn't serialize bools as numbers), but the validation gate exists for defense-in-depth so it should be actually strict. --- app/assets/services/ingest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index 5716f3d5b..3b6dc237c 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -414,8 +414,8 @@ def _backfill_image_dimensions_from_siblings( width = meta.get("width") height = meta.get("height") if ( - not isinstance(width, int) - or not isinstance(height, int) + type(width) is not int + or type(height) is not int or width <= 0 or height <= 0 ):