222 lines
7.1 KiB
Python
222 lines
7.1 KiB
Python
"""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,
|
|
)
|