unwrap-clothes/unwrap_clothes.py
2026-09-10 23:20:47 +02:00

118 lines
3 KiB
Python
Executable file

#!/usr/bin/env python3
"""Unwrinkle clothing photos and unify backgrounds for marketplace listings.
Preserves garment shape, colors, prints, and defects while removing creases
and replacing the background with a studio backdrop.
"""
import argparse
import os
import sys
from cleaner import (
API_URL,
BACKGROUND_PRESETS,
CleanerError,
CleanResult,
DEFAULT_MODEL,
EXTENSIONS,
build_prompt,
clean_garment,
encode_image,
encode_image_bytes,
restore_resolution,
restore_resolution_bytes,
)
__all__ = [
"API_URL",
"DEFAULT_MODEL",
"BACKGROUND_PRESETS",
"EXTENSIONS",
"build_prompt",
"encode_image",
"encode_image_bytes",
"restore_resolution",
"restore_resolution_bytes",
"clean_garment",
"CleanResult",
"CleanerError",
"parse_args",
"main",
]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Unwrinkle clothing photos and unify backgrounds via OpenRouter."
)
parser.add_argument("image", help="Path to input garment photo")
parser.add_argument(
"-o",
"--output",
help="Path for cleaned output image (defaults to <name>_clean.<ext>)",
)
parser.add_argument(
"--bg",
default="beige",
help="Background preset ('beige', 'white') or a custom background description",
)
parser.add_argument(
"--model",
default=DEFAULT_MODEL,
help=f"OpenRouter image model slug (default: {DEFAULT_MODEL})",
)
parser.add_argument(
"--api-key",
default=os.environ.get("OPENROUTER_API_KEY"),
help="OpenRouter API key (defaults to OPENROUTER_API_KEY env var)",
)
parser.add_argument(
"--no-restore-res",
action="store_true",
help="Skip upscaling back to source photo dimensions",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
if not os.path.isfile(args.image):
sys.exit(f"Error: file not found: {args.image}")
if not args.api_key:
sys.exit(
"Error: missing API key. Set OPENROUTER_API_KEY in your environment or pass --api-key."
)
prompt = build_prompt(args.bg)
print(f"Submitting to {args.model}...")
try:
result = clean_garment(
image_data=args.image,
prompt=prompt,
api_key=args.api_key,
model=args.model,
restore_res=not args.no_restore_res,
)
except CleanerError as e:
sys.exit(str(e))
ext = EXTENSIONS.get(result.media_type, ".jpg")
out = args.output or (os.path.splitext(args.image)[0] + "_clean" + ext)
with open(out, "wb") as f:
f.write(result.image_bytes)
if not args.no_restore_res and result.was_rescaled:
print(f"Restored resolution: {result.model_dimensions} -> {result.original_dimensions}")
cost = result.cost
cost_str = f" (cost ${cost:.2f})" if cost is not None else ""
print(f"Saved: {out}{cost_str}")
if __name__ == "__main__":
main()