diff --git a/cleaner.py b/cleaner.py index 4caac6f..b18ada3 100644 --- a/cleaner.py +++ b/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:;base64, + 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, ) diff --git a/unwrap_clothes.py b/unwrap_clothes.py index 5098ab2..f68274e 100755 --- a/unwrap_clothes.py +++ b/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