US07-03: Harden Media and Metadata Edge Cases

photo_pipeline/integrations/nsfw_model.py

@@ -15,6 +15,8 @@ from __future__ import annotations
 
 from pathlib import Path
 
+from photo_pipeline import imaging
+
 MODEL_ID = "AdamCodd/vit-base-nsfw-detector"
 BATCH = 16
 
@@ -54,13 +56,17 @@ class NsfwModel:
         self._ensure_loaded()
         import numpy as np
         import torch
-        from PIL import Image, ImageFile
+        from PIL import Image
 
-        ImageFile.LOAD_TRUNCATED_IMAGES = True
-
-        def preprocess(image):
-            image = image.convert("RGB").resize((self._size, self._size), Image.BILINEAR)
-            array = (np.asarray(image, dtype="float32") / 255.0 - 0.5) / 0.5
+        # The donor set ``ImageFile.LOAD_TRUNCATED_IMAGES = True`` here. That flag is
+        # process-global: in this application the same process also hashes files and
+        # renders previews, and those must keep failing loudly on a truncated file
+        # rather than quietly working on half of one (US07-03). An unreadable image
+        # is skipped instead — it stays unscored, and therefore visibly undecided.
+        def preprocess(path):
+            with imaging.open_image(path) as image:
+                small = image.convert("RGB").resize((self._size, self._size), Image.BILINEAR)
+            array = (np.asarray(small, dtype="float32") / 255.0 - 0.5) / 0.5
             return torch.from_numpy(array).permute(2, 0, 1)
 
         results: list[tuple[str, float]] = []
@@ -69,9 +75,10 @@ class NsfwModel:
             tensors, batch_paths = [], []
             for path in items[start : start + self.batch]:
                 try:
-                    tensors.append(preprocess(Image.open(path)))
+                    tensors.append(preprocess(path))
                     batch_paths.append(path)
-                except Exception:
+                except (imaging.MediaError, OSError, ValueError):
+                    # One bad file must not cost the batch its other fifteen.
                     continue
             if not tensors:
                 continue