mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-07 06:20:15 -03:00
feat(recipes): add get_image_dimensions helper with LRU cache
This commit is contained in:
@@ -1,9 +1,10 @@
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
from io import BytesIO
|
||||
from typing import Any, Optional
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
import piexif
|
||||
from PIL import Image, PngImagePlugin
|
||||
@@ -17,6 +18,22 @@ except ImportError:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=2048)
|
||||
def _get_image_dimensions_cached(path: str, _mtime_ns: int, _size: int) -> Optional[Tuple[int, int]]:
|
||||
"""Return ``(width, height)`` for ``path``, or ``None`` on any failure.
|
||||
|
||||
The ``_mtime_ns`` and ``_size`` arguments are part of the cache key only;
|
||||
they invalidate the entry when the file is replaced with a new image, so a
|
||||
stale preview never serves outdated dimensions.
|
||||
"""
|
||||
try:
|
||||
with Image.open(path) as img:
|
||||
return img.size
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class ExifUtils:
|
||||
"""Utility functions for working with EXIF data in images"""
|
||||
|
||||
@@ -422,6 +439,25 @@ class ExifUtils:
|
||||
# Metadata is in the middle of the string
|
||||
return user_comment[:recipe_marker_index] + user_comment[next_line_index:]
|
||||
|
||||
@staticmethod
|
||||
def get_image_dimensions(image_path: str) -> Optional[Tuple[int, int]]:
|
||||
"""Return ``(width, height)`` for an image, or ``None`` if unavailable.
|
||||
|
||||
Video containers (``.mp4``/``.webm``/``.avi``) and formats PIL cannot
|
||||
read (``.avif``/``.jxl``) return ``None`` before PIL is invoked.
|
||||
Missing or corrupt files return ``None``. Never raises.
|
||||
"""
|
||||
try:
|
||||
ext = os.path.splitext(image_path)[1].lower()
|
||||
if ext in ('.mp4', '.webm', '.avi', '.avif', '.jxl'):
|
||||
return None
|
||||
stat = os.stat(image_path)
|
||||
return _get_image_dimensions_cached(
|
||||
image_path, stat.st_mtime_ns, stat.st_size
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=False):
|
||||
"""
|
||||
|
||||
@@ -288,3 +288,58 @@ class TestIsobmffBrotliExtraction:
|
||||
# Direct extraction should return None because decompressed size exceeds limit
|
||||
result = ExifUtils._extract_isobmff_brotli(str(path))
|
||||
assert result is None
|
||||
|
||||
|
||||
# --- get_image_dimensions tests ---
|
||||
|
||||
|
||||
def test_get_image_dimensions_returns_actual_size(tmp_path):
|
||||
"""(a) A valid image returns its real (width, height)."""
|
||||
image_path = tmp_path / "preview.png"
|
||||
Image.new("RGB", (64, 32), color="red").save(image_path)
|
||||
|
||||
assert ExifUtils.get_image_dimensions(str(image_path)) == (64, 32)
|
||||
|
||||
|
||||
def test_get_image_dimensions_missing_path_returns_none(tmp_path):
|
||||
"""(b) A nonexistent path returns None without raising."""
|
||||
assert ExifUtils.get_image_dimensions(str(tmp_path / "missing.png")) is None
|
||||
|
||||
|
||||
def test_get_image_dimensions_skips_video_extension_without_pil(tmp_path, monkeypatch):
|
||||
"""(c) A .mp4 path returns None and never invokes PIL."""
|
||||
video_path = tmp_path / "preview.mp4"
|
||||
video_path.write_bytes(b"not really a video")
|
||||
|
||||
def fail_if_called(*args, **kwargs):
|
||||
raise AssertionError("PIL Image.open must not be called for video paths")
|
||||
|
||||
monkeypatch.setattr("py.utils.exif_utils.Image.open", fail_if_called)
|
||||
|
||||
assert ExifUtils.get_image_dimensions(str(video_path)) is None
|
||||
|
||||
|
||||
def test_get_image_dimensions_corrupt_file_returns_none(tmp_path):
|
||||
"""(d) A corrupt file returns None without raising."""
|
||||
image_path = tmp_path / "corrupt.png"
|
||||
image_path.write_bytes(b"\x00\x01\x02\x03 not a real image")
|
||||
|
||||
assert ExifUtils.get_image_dimensions(str(image_path)) is None
|
||||
|
||||
|
||||
def test_get_image_dimensions_cache_key_includes_mtime(tmp_path):
|
||||
"""(e) Replacing a path with a different-size image returns the new size."""
|
||||
image_path = tmp_path / "replaced.png"
|
||||
Image.new("RGB", (64, 32), color="red").save(image_path)
|
||||
assert ExifUtils.get_image_dimensions(str(image_path)) == (64, 32)
|
||||
|
||||
Image.new("RGB", (100, 50), color="blue").save(image_path)
|
||||
assert ExifUtils.get_image_dimensions(str(image_path)) == (100, 50)
|
||||
|
||||
|
||||
def test_get_image_dimensions_skips_unreadable_formats(tmp_path):
|
||||
"""(f) .avif/.jxl paths return None without raising."""
|
||||
for ext in (".avif", ".jxl"):
|
||||
image_path = tmp_path / f"preview{ext}"
|
||||
image_path.write_bytes(b"fake container data")
|
||||
assert ExifUtils.get_image_dimensions(str(image_path)) is None
|
||||
|
||||
Reference in New Issue
Block a user