import sys, os
sys.path.insert(0, "services/recolor")
import numpy as np
import cv2
from PIL import Image
from fabric_recolor.segmentation.birefnet import BirefnetMatteSegmenter
from fabric_recolor.segmentation.prompts import NEGATIVE_PROMPTS, POSITIVE_PROMPTS

OUT = "tmp/masks_v9"
os.makedirs(OUT, exist_ok=True)
seg = BirefnetMatteSegmenter()
masters = {
    "detail": "data/media/assets/20/20282cb1bd8a58596050b99b9f29e079145a7bcc45ea46b8523ae2389801fbdc.png",
    "packshot_front": "data/media/assets/79/7987376a67865864247263e5028a83950a2fd900740bbe339fab1ff8e6b3fbb2.png",
    "hero": "data/media/assets/f0/f06282d9ba621d137d7c978cf931792ffbcb6b5c9764cf94e964f83ded13ddcf.png",
}
for name, mp in masters.items():
    rgb = np.asarray(Image.open(mp).convert("RGB"), dtype=np.uint8)
    res = seg.segment(rgb, list(POSITIVE_PROMPTS), list(NEGATIVE_PROMPTS))
    alpha = res.alpha
    cov = float((alpha > 10/255).mean())
    print(f"{name}: cov={cov:.4f} notes={res.notes}")
    Image.fromarray((alpha*255).astype(np.uint8)).save(f"{OUT}/{name}.mask.png")
    ov = rgb.copy()
    edge = (alpha > 10/255) & (alpha < 200/255)
    ov[alpha >= 200/255] = (ov[alpha >= 200/255]*0.55 + np.array([255,0,0])*0.45).astype(np.uint8)
    ov[edge] = (ov[edge]*0.3 + np.array([255,255,0])*0.7).astype(np.uint8)
    Image.fromarray(ov).save(f"{OUT}/{name}.overlay.png")