"""Core garment photo cleaning logic for CLI and API reuse.""" 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 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.""" try: from PIL import Image except ImportError: print("Pillow not installed; skipping resolution restore (pip install Pillow)") return 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): """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 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_width: int | None = None orig_height: int | None = None 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") 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_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 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 ) 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 return CleanResult( image_bytes=gen_bytes, media_type=media_type, cost=cost, width=width, height=height, original_dimensions=orig_dimensions, model_dimensions=model_dimensions, )