#!/usr/bin/env python3
"""
Create OneTrainer *-masklabel.png masks from concept/dataset images
using rembg + the u2net model.

Prerequisites:
    pip install rembg pillow numpy

    The first run downloads u2net.onnx (~176 MB) automatically.

Usage:
    python3 make_masks.py /path/to/concept-folder

    Optional:
        python3 make_masks.py /path/to/concept-folder --overwrite
            regenerate masks even if they already exist

Notes:
    - Processes .png / .jpg / .jpeg / .webp
    - Skips files whose name already ends with -masklabel.png
    - Output name: <image_stem>-masklabel.png
"""

import argparse
import sys
from pathlib import Path

import numpy as np
from PIL import Image
from rembg import new_session, remove

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}


def make_mask(img_pil, session):
    result = remove(img_pil, session=session)
    alpha = np.array(result.split()[-1])
    mask = (alpha > 30).astype(np.uint8) * 255
    return Image.fromarray(mask, mode="L")


def main():
    parser = argparse.ArgumentParser(
        description="Generate OneTrainer -masklabel.png files with rembg u2net."
    )
    parser.add_argument(
        "folder",
        help="Path to the concept/dataset folder containing source images",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Regenerate masks even if they already exist",
    )
    args = parser.parse_args()

    src_dir = Path(args.folder).expanduser().resolve()
    if not src_dir.is_dir():
        print(f"Folder not found: {src_dir}")
        sys.exit(1)

    session = new_session("u2net")

    images = [
        p
        for p in sorted(src_dir.iterdir())
        if p.suffix.lower() in IMAGE_EXTS and not p.name.endswith("-masklabel.png")
    ]

    print(f"Found {len(images)} images in {src_dir}\n")

    for path in images:
        out_path = src_dir / f"{path.stem}-masklabel.png"

        if out_path.exists() and not args.overwrite:
            print(f"SKIP  {path.stem}  (mask already exists)")
            continue

        img = Image.open(path).convert("RGB")
        print(f"Masking {path.stem}...", end=" ", flush=True)
        mask = make_mask(img, session)
        mask.save(out_path)

        arr = np.array(mask)
        coverage = (arr > 127).sum() / arr.size * 100
        print(f"coverage {coverage:.1f}%")

    print("\nDone.")


if __name__ == "__main__":
    main()