mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-07 14:30: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 json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import struct
|
import struct
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional, Tuple
|
||||||
|
|
||||||
import piexif
|
import piexif
|
||||||
from PIL import Image, PngImagePlugin
|
from PIL import Image, PngImagePlugin
|
||||||
@@ -17,6 +18,22 @@ except ImportError:
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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:
|
class ExifUtils:
|
||||||
"""Utility functions for working with EXIF data in images"""
|
"""Utility functions for working with EXIF data in images"""
|
||||||
|
|
||||||
@@ -422,6 +439,25 @@ class ExifUtils:
|
|||||||
# Metadata is in the middle of the string
|
# Metadata is in the middle of the string
|
||||||
return user_comment[:recipe_marker_index] + user_comment[next_line_index:]
|
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
|
@staticmethod
|
||||||
def optimize_image(image_data, target_width=250, format='webp', quality=85, preserve_metadata=False):
|
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
|
# Direct extraction should return None because decompressed size exceeds limit
|
||||||
result = ExifUtils._extract_isobmff_brotli(str(path))
|
result = ExifUtils._extract_isobmff_brotli(str(path))
|
||||||
assert result is None
|
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