Merge pull request #3001 from AL3708/perf/image-metadata-only-scan

Extend metadata-only bucketing to still images in the discovery backend

Author: bghira

simpletuner/helpers/metadata/backends/discovery.py

@@ -6,6 +6,8 @@
 from io import BytesIO
 from typing import Optional
 
+from PIL import Image
+
 from simpletuner.helpers.data_backend.base import BaseDataBackend
 from simpletuner.helpers.data_backend.dataset_types import DatasetType
 from simpletuner.helpers.image_manipulation.load import load_image, load_video
@@ -94,16 +96,58 @@ def __init__(
             max_num_samples=max_num_samples,
         )
 
-    def _should_use_metadata_only_for_video(self) -> bool:
-        if not _is_ffprobe_available():
-            return False
-        if self.dataset_type is not DatasetType.VIDEO:
-            return False
+    def _should_use_metadata_only(self, is_video_file: bool) -> bool:
+        """Whether this sample can be bucketed from dimensions alone.
+
+        Bucketing needs only the sample dimensions: calculate_target_size() and
+        the crop coordinate maths derive everything from original_size, and no
+        pixel-derived metadata is stored. Videos already take this path via
+        ffprobe; still images can do the same from the file header.
+
+        Face cropping is the one crop style whose coordinates depend on pixel
+        content, so it always requires a full decode.
+        """
         crop_enabled = bool(self.dataset_config.get("crop", False))
         crop_style = str(self.dataset_config.get("crop_style") or "random").lower()
         if crop_enabled and crop_style == "face":
             return False
-        return True
+
+        if is_video_file:
+            return self.dataset_type is DatasetType.VIDEO and _is_ffprobe_available()
+        if self.dataset_type in (DatasetType.IMAGE, DatasetType.CONDITIONING):
+            # Header reads need a real path; remote backends stay on full reads.
+            return getattr(self.data_backend, "type", None) == "local"
+        return False
+
+    def _probe_image_dimensions(self, image_path_str: str) -> Optional[dict]:
+        """Read image dimensions from the file header, without decoding pixels.
+
+        PIL parses only the header on Image.open(); pixel data is loaded lazily
+        and never requested here. Returns None to signal that the caller should
+        fall back to a full decode.
+
+        EXIF orientations 5-8 transpose the image, so the stored dimensions are
+        swapped relative to what exif_transpose() produces on the full-decode
+        path. The orientation tag lives in the header too, so it is applied here
+        to keep both paths in agreement.
+        """
+        try:
+            with Image.open(image_path_str) as image:
+                width, height = image.size
+                orientation = image.getexif().get(0x0112)
+            if orientation in (5, 6, 7, 8):
+                width, height = height, width
+        except Exception as e:
+            logger.debug(
+                "(id=%s) Could not read image header for %s (%s); falling back to full decode.",
+                self.id,
+                image_path_str,
+                e,
+            )
+            return None
+        if not width or not height:
+            return None
+        return {"original_size": (width, height)}
 
     def _needs_video_frame_count(self) -> bool:
         if self.bucket_strategy == "resolution_frames":
@@ -382,7 +426,7 @@ def _process_for_bucket(
             is_video_file = file_extension.strip(".") in video_file_extensions
 
             use_metadata_only = False
-            if is_video_file and self._should_use_metadata_only_for_video():
+            if is_video_file and self._should_use_metadata_only(is_video_file):
                 if getattr(self.data_backend, "type", None) == "local":
                     video_metadata = self._probe_video_metadata(image_path_str, None)
                 else:
@@ -429,6 +473,13 @@ def _process_for_bucket(
 
                     logger.debug("(id=%s) Using ffprobe metadata-only scan for %s", self.id, image_path_str)
 
+            elif not is_video_file and self._should_use_metadata_only(is_video_file):
+                image_dimensions = self._probe_image_dimensions(image_path_str)
+                if image_dimensions:
+                    use_metadata_only = True
+                    image_metadata.update(image_dimensions)
+                    logger.debug("(id=%s) Using header-only metadata scan for %s", self.id, image_path_str)
+
             if use_metadata_only:
                 if not self.meets_resolution_requirements(image_metadata=image_metadata):
                     if not self.delete_unwanted_images:

tests/test_metadata_backend.py

@@ -63,6 +63,76 @@ def test_len(self):
         }
         self.assertEqual(len(self.metadata_backend), 3)
 
+    def test_probe_image_dimensions_applies_exif_orientation(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            image_path = Path(tmpdir) / "oriented.jpg"
+            exif = Image.Exif()
+            exif[0x0112] = 6
+            Image.new("RGB", (64, 32)).save(image_path, exif=exif)
+
+            metadata = self.metadata_backend._probe_image_dimensions(str(image_path))
+
+        self.assertEqual(metadata, {"original_size": (32, 64)})
+
+    def test_local_image_metadata_scan_skips_backend_read(self):
+        self.data_backend.type = "local"
+        self.metadata_backend.dataset_type = DatasetType.IMAGE
+        self.metadata_backend.dataset_config = {"dataset_type": "image", "crop": False}
+        self.data_backend.read.reset_mock()
+        prepared = SimpleNamespace(
+            crop_coordinates=(0, 0),
+            target_size=(512, 256),
+            intermediary_size=(512, 256),
+            aspect_ratio=2.0,
+        )
+
+        with (
+            patch.object(
+                self.metadata_backend,
+                "_probe_image_dimensions",
+                return_value={"original_size": (512, 256)},
+            ),
+            patch("simpletuner.helpers.metadata.backends.discovery.TrainingSample") as training_sample,
+        ):
+            training_sample.return_value.prepare.return_value = prepared
+            metadata_updates = {}
+            buckets = self.metadata_backend._process_for_bucket(
+                "image.png",
+                {},
+                metadata_updates=metadata_updates,
+            )
+
+        self.data_backend.read.assert_not_called()
+        self.assertEqual(buckets, {"2.0": ["image.png"]})
+        self.assertEqual(metadata_updates["image.png"]["original_size"], (512, 256))
+
+    def test_face_crop_forces_full_image_decode(self):
+        self.data_backend.type = "local"
+        self.metadata_backend.dataset_type = DatasetType.IMAGE
+        self.metadata_backend.dataset_config = {
+            "dataset_type": "image",
+            "crop": True,
+            "crop_style": "face",
+        }
+        self.data_backend.read = Mock(return_value=b"image payload")
+        prepared = SimpleNamespace(
+            crop_coordinates=(0, 0),
+            target_size=(512, 256),
+            intermediary_size=(512, 256),
+            aspect_ratio=2.0,
+        )
+
+        with (
+            patch.object(self.metadata_backend, "_probe_image_dimensions") as probe_dimensions,
+            patch("simpletuner.helpers.metadata.backends.discovery.load_image", return_value=self.test_image),
+            patch("simpletuner.helpers.metadata.backends.discovery.TrainingSample") as training_sample,
+        ):
+            training_sample.return_value.prepare.return_value = prepared
+            self.metadata_backend._process_for_bucket("image.png", {})
+
+        probe_dimensions.assert_not_called()
+        self.data_backend.read.assert_called_once_with("image.png")
+
     def test_discover_new_files(self):
         # Assuming that StateTracker.get_image_files returns known files
         # and list_files should return both known and potentially new files