pak-jabit/catalogue/photoroom_pipeline.py

#!/usr/bin/env python3
"""
Pak Jabit — Product Image Pipeline
PhotoRoom API: remove background + studio-quality packshot

Usage:
  python3 photoroom_pipeline.py --input products/01-gshock-*/image-1.jpg
  python3 photoroom_pipeline.py --input-dir products/ --all
  python3 photoroom_pipeline.py --input products/01-gshock-*/image-1.jpg --dry-run

Env:
  PHOTOROOM_API_KEY  — sandbox_ prefix for testing (watermarked), real key for production
"""

import argparse, glob, json, os, sys, time
from pathlib import Path
import urllib.request, urllib.error

API_KEY = os.environ.get("PHOTOROOM_API_KEY")
EDIT_URL = "https://image-api.photoroom.com/v2/edit"
SEGMENT_URL = "https://sdk.photoroom.com/v1/segment"

def check_account():
    req = urllib.request.Request(
        "https://image-api.photoroom.com/v2/account",
        headers={"x-api-key": API_KEY},
    )
    with urllib.request.urlopen(req, timeout=15) as r:
        data = json.loads(r.read())
    print(f"Plan: {data['plan']} | Credits: {data['images']['available']}/{data['images']['subscription']}")
    return data

def remove_background(image_path: str, output_path: str):
    boundary = "----PakJabitBoundary"
    filename = os.path.basename(image_path)
    with open(image_path, "rb") as f:
        image_data = f.read()

    body = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="image_file"; filename="{filename}"\r\n'
        f"Content-Type: image/jpeg\r\n\r\n"
    ).encode() + image_data + (
        f"\r\n--{boundary}\r\n"
        f'Content-Disposition: form-data; name="format"\r\n\r\npng\r\n'
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="size"\r\n\r\nhd\r\n'
        f"--{boundary}--\r\n"
    ).encode()

    req = urllib.request.Request(SEGMENT_URL, data=body, method="POST", headers={
        "x-api-key": API_KEY,
        "Content-Type": f"multipart/form-data; boundary={boundary}",
    })
    with urllib.request.urlopen(req, timeout=60) as r:
        with open(output_path, "wb") as out:
            out.write(r.read())
    return output_path

def studio_packshot(image_path: str, output_path: str, bg_color: str = "FFFFFF"):
    boundary = "----PakJabitBoundary"
    filename = os.path.basename(image_path)
    with open(image_path, "rb") as f:
        image_data = f.read()

    fields = [
        ("background.color", bg_color),
        ("padding", "0.1"),
        ("shadow.mode", "ai.soft"),
        ("outputSize", "1200x1600"),
        ("export.format", "jpeg"),
        ("removeBackground", "true"),
    ]

    body = b""
    for name, value in fields:
        body += (
            f"--{boundary}\r\n"
            f'Content-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n'
        ).encode()

    body += (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="imageFile"; filename="{filename}"\r\n'
        f"Content-Type: image/jpeg\r\n\r\n"
    ).encode() + image_data + f"\r\n--{boundary}--\r\n".encode()

    req = urllib.request.Request(EDIT_URL, data=body, method="POST", headers={
        "x-api-key": API_KEY,
        "Content-Type": f"multipart/form-data; boundary={boundary}",
    })
    with urllib.request.urlopen(req, timeout=90) as r:
        with open(output_path, "wb") as out:
            out.write(r.read())
    return output_path

def process_image(image_path: str, output_dir: str, mode: str = "packshot"):
    stem = Path(image_path).stem
    parent = Path(image_path).parent.name

    out_subdir = os.path.join(output_dir, parent)
    os.makedirs(out_subdir, exist_ok=True)

    if mode == "remove-bg":
        out_path = os.path.join(out_subdir, f"{stem}-nobg.png")
        remove_background(image_path, out_path)
    elif mode == "packshot":
        out_path = os.path.join(out_subdir, f"{stem}-studio.jpg")
        studio_packshot(image_path, out_path)
    else:
        raise ValueError(f"Unknown mode: {mode}")

    print(f"  ✅ {os.path.basename(image_path)} → {os.path.basename(out_path)}")
    return out_path

def main():
    parser = argparse.ArgumentParser(description="PhotoRoom product image pipeline")
    parser.add_argument("--input", nargs="+", help="Image file(s) or glob pattern")
    parser.add_argument("--input-dir", help="Process all images in product subdirs")
    parser.add_argument("--all", action="store_true", help="Process all image-*.jpg in input-dir")
    parser.add_argument("--first-only", action="store_true", help="Only process image-1.jpg per product")
    parser.add_argument("--output-dir", default="products-enhanced", help="Output directory")
    parser.add_argument("--mode", choices=["remove-bg", "packshot"], default="packshot")
    parser.add_argument("--dry-run", action="store_true", help="List files without processing")
    parser.add_argument("--account", action="store_true", help="Check account credits")
    parser.add_argument("--delay", type=float, default=1.0, help="Delay between API calls (seconds)")
    args = parser.parse_args()

    if not API_KEY:
        print("❌ Set PHOTOROOM_API_KEY env var first")
        print("   Sandbox key (free, watermarked): starts with sandbox_")
        print("   Get it from: https://app.photoroom.com/api-dashboard")
        sys.exit(1)

    if args.account:
        check_account()
        return

    images = []
    if args.input:
        for pattern in args.input:
            images.extend(glob.glob(pattern))
    elif args.input_dir and args.all:
        pattern = "image-1.jpg" if args.first_only else "image-*.jpg"
        images = sorted(glob.glob(os.path.join(args.input_dir, "**", pattern), recursive=True))
    else:
        parser.print_help()
        return

    if not images:
        print("No images found")
        return

    print(f"📸 {len(images)} images | mode: {args.mode} | output: {args.output_dir}")

    if args.dry_run:
        for img in images:
            print(f"  would process: {img}")
        return

    check_account()
    print()

    ok, fail = 0, 0
    for i, img in enumerate(images, 1):
        print(f"[{i}/{len(images)}] {img}")
        try:
            process_image(img, args.output_dir, args.mode)
            ok += 1
        except urllib.error.HTTPError as e:
            print(f"  ❌ HTTP {e.code}: {e.read().decode()[:200]}")
            fail += 1
            if e.code == 429:
                print("  ⏳ Rate limited, waiting 30s...")
                time.sleep(30)
        except Exception as e:
            print(f"  ❌ {e}")
            fail += 1

        if i < len(images):
            time.sleep(args.delay)

    print(f"\n✅ Done: {ok} ok, {fail} failed")

if __name__ == "__main__":
    main()