250 lines
8.3 KiB
Python
250 lines
8.3 KiB
Python
"""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 = os.environ.get("OPENROUTER_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 resolve_api_key(api_key: str | None = None) -> str | None:
|
|
"""Resolve OpenRouter API key from parameter or OPENROUTER_API_KEY environment variable."""
|
|
if api_key and api_key.strip():
|
|
return api_key.strip()
|
|
return os.environ.get("OPENROUTER_API_KEY")
|
|
|
|
|
|
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:"):
|
|
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
|
|
|
|
|
|
_get_image_info = get_image_info
|
|
|
|
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 to_payload(self) -> dict[str, Any]:
|
|
"""Serialize result to a dictionary suitable for API responses."""
|
|
return {
|
|
"image": encode_image_bytes(self.image_bytes, mime=self.media_type),
|
|
"cost": self.cost,
|
|
"width": self.width,
|
|
"height": self.height,
|
|
"original_dimensions": list(self.original_dimensions)
|
|
if self.original_dimensions
|
|
else None,
|
|
"model_dimensions": list(self.model_dimensions)
|
|
if self.model_dimensions
|
|
else None,
|
|
"was_rescaled": self.was_rescaled,
|
|
"media_type": self.media_type,
|
|
}
|
|
|
|
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,
|
|
)
|