118 lines
3 KiB
Python
Executable file
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()
|