`enrich-palette.py`

40 lines
#!/usr/bin/env python3
"""Ground-truth palette + a NORMALIZED image for the VL.
- palette: exact hex + area-% from real pixels (more accurate than any VLM), near-dup merged.
- image_b64: the image re-encoded as a clean RGB JPEG, long-edge-capped. Normalizing through
  Pillow guarantees ollama/llama-server can load it (catalog images are often CMYK JPEG / WebP /
  odd encodings that Pillow reads but the VL image loader rejects → "Failed to load image").
Usage: enrich-palette.py <image> [k]   → prints {"palette":[...], "image_b64":"..."}"""
import sys, json, io, base64, os
from PIL import Image

def _dist(a, b):
    return sum((x - y) ** 2 for x, y in zip(a, b)) ** 0.5

def analyze(path, k=6, merge_thresh=30, vl_max=int(os.environ.get("ENRICH_VL_MAX", "1024"))):
    im = Image.open(path).convert("RGB")            # CMYK/RGBA/P/WebP -> RGB
    # --- palette from a small copy ---
    small = im.copy(); small.thumbnail((220, 220))
    q = small.quantize(colors=16, method=Image.MEDIANCUT)
    pal = q.getpalette(); counts = sorted(q.getcolors(), reverse=True)
    total = sum(c for c, _ in counts) or 1
    buckets = []
    for count, idx in counts:
        rgb = tuple(pal[idx*3:idx*3+3]); pct = count / total * 100
        for b in buckets:
            if _dist(rgb, b["rgb"]) <= merge_thresh:
                b["pct"] += pct; break
        else:
            buckets.append({"rgb": rgb, "pct": pct})
    buckets = [b for b in buckets if b["pct"] >= 2][:k]
    out = [{"hex": "#%02X%02X%02X" % b["rgb"], "percentage": round(b["pct"])} for b in buckets]
    diff = 100 - sum(c["percentage"] for c in out)
    if out: out[0]["percentage"] += diff
    # --- normalized VL image (clean JPEG ollama always loads) ---
    vl = im.copy(); vl.thumbnail((vl_max, vl_max))
    buf = io.BytesIO(); vl.save(buf, format="JPEG", quality=90)
    return {"palette": out, "image_b64": base64.b64encode(buf.getvalue()).decode()}

if __name__ == "__main__":
    print(json.dumps(analyze(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 6)))