"""Core garment photo cleaning logic for CLI and API reuse.""" from __future__ import annotations import base64 from dataclasses import dataclass import io import mimetypes import os from typing import Any import requests API_URL = "https://openrouter.ai/api/v1/images" DEFAULT_MODEL = os.environ.get("MODEL", "meta/muse-image") BACKGROUND_PRESETS = { "beige": "Soft warm beige studio background, evenly lit, no shadows, no props.", "white": "Plain bright white studio backdrop, evenly lit, no shadows, no gradients, no props.", } EXTENSIONS = { "image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp", } 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 ( "Product photography edit of this garment photo. Keep the garments pixel-faithful: identical " "shape, colors, prints, labels, buttons, and every existing flaw unchanged. Remove ALL wrinkles " "and creases completely: the fabric must look perfectly smooth and freshly ironed, flat like a " f"new catalog product photo. {bg_description}" ) def encode_image_bytes(data: bytes, mime: str = "image/jpeg") -> str: return f"data:{mime};base64,{base64.b64encode(data).decode()}" def encode_image(path: str) -> str: mime = mimetypes.guess_type(path)[0] or "image/png" with open(path, "rb") as f: return encode_image_bytes(f.read(), mime=mime) def restore_resolution_bytes( gen_bytes: bytes, orig_width: int, orig_height: int, fmt: str = "JPEG", ) -> bytes: """Resize generated image bytes back to the original dimensions using Lanczos resampling.""" try: from PIL import Image except ImportError: return gen_bytes with Image.open(io.BytesIO(gen_bytes)) as gen: if (gen.width, gen.height) == (orig_width, orig_height): return gen_bytes resample = getattr(Image, "Resampling", Image).LANCZOS resized = gen.resize((orig_width, orig_height), resample) buf = io.BytesIO() kwargs = {"quality": 95} if fmt.upper() in ("JPEG", "JPG") else {} resized.save(buf, format=fmt, **kwargs) return buf.getvalue() def restore_resolution(out: str, src: str) -> None: """Resize the generated image back to the input dimensions using Lanczos resampling.""" 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}") class CleanerError(Exception): """Base error raised when cleaning or calling OpenRouter fails.""" pass @dataclass class CleanResult: image_bytes: bytes media_type: str = "image/jpeg" cost: float | None = None width: int | None = None height: int | None = None 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, prompt: str, api_key: str, model: str = DEFAULT_MODEL, api_url: str = API_URL, timeout: int = 300, ) -> dict: """Submit image edit request to OpenRouter images API.""" headers = {"Authorization": f"Bearer {api_key}"} payload: dict[str, Any] = { "model": model, "prompt": prompt, "background": "opaque", "input_references": [ {"type": "image_url", "image_url": {"url": image_url_or_data_uri}} ], "output_format": "jpeg", } resp = requests.post(api_url, headers=headers, json=payload, timeout=timeout) if resp.status_code != 200: try: message = resp.json()["error"]["message"] except Exception: message = resp.text raise CleanerError(f"API error {resp.status_code}: {message}") return resp.json() def clean_garment( image_data: bytes | str, prompt: str, api_key: str, model: str = DEFAULT_MODEL, restore_res: bool = True, mime: str = "image/jpeg", api_url: str = API_URL, timeout: int = 300, ) -> CleanResult: """Clean garment image given file path or bytes, returning CleanResult.""" 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): 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) else: raise TypeError("image_data must be bytes or str path") result = call_openrouter_images( image_url_or_data_uri=data_uri, prompt=prompt, api_key=api_key, model=model, api_url=api_url, timeout=timeout, ) image_entry = result["data"][0] gen_bytes = base64.b64decode(image_entry["b64_json"]) media_type = image_entry.get("media_type", "image/jpeg") cost = result.get("usage", {}).get("cost") 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" gen_bytes = restore_resolution_bytes( gen_bytes, orig_width, orig_height, fmt=fmt ) final_dims, _ = _get_image_info(gen_bytes) width, height = final_dims if final_dims else (None, None) return CleanResult( image_bytes=gen_bytes, media_type=media_type, cost=cost, width=width, height=height, original_dimensions=orig_dims, model_dimensions=model_dimensions, )