#!/usr/bin/env python3
"""
remove_bg.py — batch background removal, free and offline after first run.

Point it at one image or a whole folder; it saves transparent PNGs next to
(or into) an output folder. Uses rembg (U^2-Net under the hood) — MIT
licensed, no API key, no per-image cost, no upload to a third-party service.

Usage:
    python remove_bg.py photo.jpg
    python remove_bg.py ./product-shots/ -o ./product-shots/no-bg
    python remove_bg.py ./product-shots/ --model isnet-general-use
"""
import argparse
import sys
from pathlib import Path

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}


def main():
    parser = argparse.ArgumentParser(description=__doc__.strip().splitlines()[0])
    parser.add_argument("input", help="Image file or folder of images")
    parser.add_argument("-o", "--output", default="no-bg", help="Output folder (default: ./no-bg)")
    parser.add_argument(
        "--model",
        default="u2net",
        help="rembg model: u2net (general, default), u2net_human_seg (people), "
        "isnet-general-use (sharper edges, slower)",
    )
    args = parser.parse_args()

    try:
        from rembg import remove, new_session
        from PIL import Image
    except ImportError:
        print("Missing dependencies. Run: pip install rembg pillow onnxruntime")
        sys.exit(1)

    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Not found: {input_path}")
        sys.exit(1)

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    if input_path.is_dir():
        images = sorted(p for p in input_path.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS)
    else:
        images = [input_path]

    if not images:
        print(f"No images found in {input_path}")
        sys.exit(1)

    print(f"Loading model '{args.model}' (first run downloads it, then it's cached)...")
    session = new_session(args.model)
    print(f"Removing backgrounds from {len(images)} image(s)...\n")

    for i, img_path in enumerate(images, 1):
        out_path = output_dir / f"{img_path.stem}.png"
        print(f"  [{i}/{len(images)}] {img_path.name} -> {out_path}", end="  ", flush=True)
        with Image.open(img_path) as img:
            result = remove(img, session=session)
            result.save(out_path)
        print("done")

    print(f"\nAll set — {len(images)} transparent PNG(s) saved to {output_dir}/")


if __name__ == "__main__":
    main()