refactor: address code review feedback on image info extraction and rescaling

cleaner.py

@@ -1,5 +1,7 @@
 """Core garment photo cleaning logic for CLI and API reuse."""
 
+from __future__ import annotations
+
 import base64
 from dataclasses import dataclass
 import io
@@ -24,6 +26,30 @@ EXTENSIONS = {
 }
 
 
+def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, str | None]:
+    """Extract (width, height) dimensions and mime type from bytes, file path, or data URI."""
+    try:
+        from PIL import Image
+        raw_bytes: bytes | None = None
+        if isinstance(data_or_path, bytes):
+            raw_bytes = data_or_path
+        elif isinstance(data_or_path, str) and data_or_path.startswith("data:"):
+            # Parse data URI: data:<mime>;base64,<payload>
+            comma_idx = data_or_path.find(",")
+            if comma_idx != -1:
+                raw_bytes = base64.b64decode(data_or_path[comma_idx + 1:])
+        elif isinstance(data_or_path, str) and os.path.isfile(data_or_path):
+            with open(data_or_path, "rb") as f:
+                raw_bytes = f.read()
+
+        if raw_bytes is not None:
+            with Image.open(io.BytesIO(raw_bytes)) as im:
+                mime = Image.MIME.get(im.format, "image/jpeg") if im.format else "image/jpeg"
+                return im.size, mime
+    except Exception:
+        pass
+    return None, None
+
 def build_prompt(bg_choice: str) -> str:
     bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice)
     return (
@@ -69,25 +95,19 @@ def restore_resolution_bytes(
 
 def restore_resolution(out: str, src: str) -> None:
     """Resize the generated image back to the input dimensions using Lanczos resampling."""
-    try:
-        from PIL import Image
-    except ImportError:
-        print("Pillow not installed; skipping resolution restore (pip install Pillow)")
+    orig_dims, _ = _get_image_info(src)
+    if not orig_dims:
         return
+    with open(out, "rb") as f:
+        out_bytes = f.read()
+    ext = os.path.splitext(out)[1].lower()
+    fmt = "JPEG" if ext in (".jpg", ".jpeg") else "PNG"
+    resized_bytes = restore_resolution_bytes(out_bytes, orig_dims[0], orig_dims[1], fmt=fmt)
+    if len(resized_bytes) != len(out_bytes):
+        with open(out, "wb") as f:
+            f.write(resized_bytes)
+        print(f"Restored resolution to {orig_dims}")
 
-    with Image.open(src) as orig, Image.open(out) as gen:
-        if gen.size == orig.size:
-            return
-        resample = getattr(Image, "Resampling", Image).LANCZOS
-        resized = gen.resize(orig.size, resample)
-        ext = os.path.splitext(out)[1].lower()
-        fmt = "JPEG" if ext in (".jpg", ".jpeg") else None
-        kwargs = {"quality": 95} if fmt == "JPEG" else {}
-        if fmt:
-            resized.save(out, fmt, **kwargs)
-        else:
-            resized.save(out)
-        print(f"Restored resolution: {gen.size} -> {orig.size}")
 
 
 class CleanerError(Exception):
@@ -105,6 +125,13 @@ class CleanResult:
     original_dimensions: tuple[int, int] | None = None
     model_dimensions: tuple[int, int] | None = None
 
+    @property
+    def was_rescaled(self) -> bool:
+        """True if resolution restoration rescaled the image from the model's output size."""
+        if self.model_dimensions and self.original_dimensions:
+            return self.model_dimensions != self.original_dimensions
+        return False
+
 
 def call_openrouter_images(
     image_url_or_data_uri: str,
@@ -146,30 +173,18 @@ def clean_garment(
     timeout: int = 300,
 ) -> CleanResult:
     """Clean garment image given file path or bytes, returning CleanResult."""
-    orig_width: int | None = None
-    orig_height: int | None = None
+    orig_dims, detected_mime = _get_image_info(image_data)
+    orig_width, orig_height = orig_dims if orig_dims else (None, None)
+    if detected_mime:
+        mime = detected_mime
 
     if isinstance(image_data, bytes):
-        try:
-            from PIL import Image
-            with Image.open(io.BytesIO(image_data)) as im:
-                orig_width, orig_height = im.size
-                if im.format:
-                    mime = Image.MIME.get(im.format, mime)
-        except Exception:
-            pass
         data_uri = encode_image_bytes(image_data, mime=mime)
     elif isinstance(image_data, str):
         if image_data.startswith("data:"):
             data_uri = image_data
         else:
             data_uri = encode_image(image_data)
-            try:
-                from PIL import Image
-                with Image.open(image_data) as im:
-                    orig_width, orig_height = im.size
-            except Exception:
-                pass
     else:
         raise TypeError("image_data must be bytes or str path")
 
@@ -186,13 +201,9 @@ def clean_garment(
     gen_bytes = base64.b64decode(image_entry["b64_json"])
     media_type = image_entry.get("media_type", "image/jpeg")
     cost = result.get("usage", {}).get("cost")
-    model_dimensions: tuple[int, int] | None = None
-    try:
-        from PIL import Image
-        with Image.open(io.BytesIO(gen_bytes)) as im:
-            model_dimensions = im.size
-    except Exception:
-        pass
+
+    model_dims, _ = _get_image_info(gen_bytes)
+    model_dimensions = model_dims
 
     if restore_res and orig_width is not None and orig_height is not None:
         fmt = "JPEG" if media_type == "image/jpeg" else "PNG"
@@ -200,16 +211,8 @@ def clean_garment(
             gen_bytes, orig_width, orig_height, fmt=fmt
         )
 
-    width: int | None = orig_width if (restore_res and orig_width is not None) else (model_dimensions[0] if model_dimensions else None)
-    height: int | None = orig_height if (restore_res and orig_height is not None) else (model_dimensions[1] if model_dimensions else None)
-    try:
-        from PIL import Image
-        with Image.open(io.BytesIO(gen_bytes)) as im:
-            width, height = im.size
-    except Exception:
-        pass
-
-    orig_dimensions = (orig_width, orig_height) if (orig_width is not None and orig_height is not None) else None
+    final_dims, _ = _get_image_info(gen_bytes)
+    width, height = final_dims if final_dims else (None, None)
 
     return CleanResult(
         image_bytes=gen_bytes,
@@ -217,6 +220,6 @@ def clean_garment(
         cost=cost,
         width=width,
         height=height,
-        original_dimensions=orig_dimensions,
+        original_dimensions=orig_dims,
         model_dimensions=model_dimensions,
     )

unwrap_clothes.py

@@ -106,12 +106,7 @@ def main() -> None:
     with open(out, "wb") as f:
         f.write(result.image_bytes)
 
-    if (
-        not args.no_restore_res
-        and result.model_dimensions is not None
-        and result.original_dimensions is not None
-        and result.model_dimensions != result.original_dimensions
-    ):
+    if not args.no_restore_res and result.was_rescaled:
         print(f"Restored resolution: {result.model_dimensions} -> {result.original_dimensions}")
 
     cost = result.cost