`scripts/vision_gate.py`

182 lines
#!/usr/bin/env python3
"""
Gemini-vision YES/NO gate for the non-PR Tier-1 color-image defect set (DTD Q1=b).
For each row: fetch image_url, ask Gemini 2.0 Flash whether the wallcovering's
DOMINANT color is consistent with the declared colorway name, or clearly different.
Drops family-parser / busy-pattern FALSE POSITIVES; keeps real-defect candidates.

Read-only against catalog (only GETs the public Shopify image + calls Gemini).
Writes results to out/vision_gate_results.jsonl + out/vision_gate_summary.csv.
"""
import csv, json, os, sys, time, base64, urllib.request, urllib.error, re

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CSV_IN = os.path.join(ROOT, "out", "tier1_confirmed_defects.csv")
JSONL_OUT = os.path.join(ROOT, "out", "vision_gate_results.jsonl")
CSV_OUT = os.path.join(ROOT, "out", "vision_gate_summary.csv")
V2_JSONL = os.path.join(ROOT, "full", "reextracted_full.jsonl")

MODEL = "gemini-2.5-flash-lite"
COST_PER_CALL = 0.0006  # ~$ per product image (always-show-costs rule)

def get_key():
    env = os.path.expanduser("~/Projects/secrets-manager/.env")
    with open(env) as f:
        for line in f:
            if line.startswith("GEMINI_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    sys.exit("no GEMINI_API_KEY")

API_KEY = get_key()
URL = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={API_KEY}"

def load_v2_verdicts():
    """Map image_url -> v2 verdict so we can tag survivors-of-v2 vs cleared."""
    m = {}
    if not os.path.exists(V2_JSONL):
        return m
    with open(V2_JSONL) as f:
        for line in f:
            try:
                r = json.loads(line)
                m[r.get("image_url")] = r.get("verdict")
            except Exception:
                pass
    return m

def fetch_image(url, tries=3):
    for i in range(tries):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=30) as r:
                data = r.read()
                ct = r.headers.get("Content-Type", "image/jpeg")
                if "image" not in ct:
                    ct = "image/jpeg"
                return data, ct
        except Exception as e:
            if i == tries - 1:
                return None, str(e)
            time.sleep(2)
    return None, "exhausted"

PROMPT_TMPL = (
    "This is a product image for a wallcovering/fabric that is merchandised online "
    "under the color name \"{declared}\". Look ONLY at the actual wallcovering/fabric "
    "surface (ignore any white studio background or border). "
    "Is the DOMINANT color of the wallcovering consistent with the name \"{declared}\", "
    "or is the image CLEARLY a different color than that name implies? "
    "A multi-color or busy pattern where \"{declared}\" is plausibly one of the inks "
    "counts as CONSISTENT (matches_declared=true). Only answer false when the image is "
    "unmistakably the wrong color for that name. "
    "Respond with STRICT JSON only, no prose: "
    '{{"dominant_color":"<plain english>","matches_declared":true|false,'
    '"confidence":<0.0-1.0>,"note":"<=12 words"}}'
)

def ask_gemini(img_bytes, ct, declared):
    prompt = PROMPT_TMPL.format(declared=declared)
    body = {
        "contents": [{
            "parts": [
                {"text": prompt},
                {"inline_data": {"mime_type": ct, "data": base64.b64encode(img_bytes).decode()}},
            ]
        }],
        "generationConfig": {"temperature": 0, "maxOutputTokens": 200},
    }
    req = urllib.request.Request(
        URL, data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    for i in range(3):
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                resp = json.loads(r.read())
            txt = resp["candidates"][0]["content"]["parts"][0]["text"]
            mt = re.search(r"\{.*\}", txt, re.S)
            if mt:
                return json.loads(mt.group(0)), None
            return None, "no-json:" + txt[:120]
        except urllib.error.HTTPError as e:
            err = e.read().decode()[:200]
            if e.code == 429:
                time.sleep(15)
                continue
            if i == 2:
                return None, f"http{e.code}:{err}"
            time.sleep(3)
        except Exception as e:
            if i == 2:
                return None, str(e)[:160]
            time.sleep(3)
    return None, "exhausted"

def main():
    v2 = load_v2_verdicts()
    rows = []
    with open(CSV_IN) as f:
        for r in csv.DictReader(f):
            if r["vendor"] != "Phillipe Romano":
                rows.append(r)
    print(f"non-PR Tier-1 rows: {len(rows)}", file=sys.stderr)

    results = []
    calls = 0
    for i, r in enumerate(rows, 1):
        declared = r["declared_color"]
        img_bytes, ct = fetch_image(r["image_url"])
        rec = {
            "vendor": r["vendor"], "dw_sku": r["dw_sku"], "mfr_sku": r["mfr_sku"],
            "declared_color": declared, "declared_family": r["declared_family"],
            "image_family": r["real_image_family"], "image_hex": r["real_image_hex"],
            "image_url": r["image_url"], "handle": r["handle"],
            "v2_verdict": v2.get(r["image_url"]),
        }
        if img_bytes is None:
            rec.update({"gate": "ERROR", "error": "fetch:" + str(ct)})
            results.append(rec)
            print(f"[{i}/{len(rows)}] FETCH-ERR {r['vendor']} {r['dw_sku']}", file=sys.stderr)
            continue
        out, err = ask_gemini(img_bytes, ct, declared)
        calls += 1
        if out is None:
            rec.update({"gate": "ERROR", "error": err})
        else:
            matches = bool(out.get("matches_declared"))
            rec.update({
                "gate": "FALSE_POSITIVE" if matches else "SURVIVOR",
                "vision_dominant": out.get("dominant_color"),
                "vision_matches": matches,
                "vision_confidence": out.get("confidence"),
                "vision_note": out.get("note"),
            })
        results.append(rec)
        g = rec.get("gate")
        print(f"[{i}/{len(rows)}] {g:14s} {r['vendor'][:18]:18s} {r['dw_sku']:14s} "
              f"declared='{declared[:30]}' vision='{rec.get('vision_dominant','')}'", file=sys.stderr)
        time.sleep(0.3)

    with open(JSONL_OUT, "w") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")
    fields = ["vendor", "dw_sku", "mfr_sku", "declared_color", "declared_family",
              "image_family", "v2_verdict", "gate", "vision_dominant",
              "vision_matches", "vision_confidence", "vision_note", "image_url", "handle"]
    with open(CSV_OUT, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        w.writeheader()
        for r in results:
            w.writerow(r)

    surv = [r for r in results if r.get("gate") == "SURVIVOR"]
    fp = [r for r in results if r.get("gate") == "FALSE_POSITIVE"]
    errs = [r for r in results if r.get("gate") == "ERROR"]
    print("\n==== VISION GATE SUMMARY ====", file=sys.stderr)
    print(f"total={len(results)} SURVIVOR={len(surv)} FALSE_POSITIVE={len(fp)} ERROR={len(errs)}", file=sys.stderr)
    print(f"gemini calls={calls}  est cost=${calls*COST_PER_CALL:.4f}", file=sys.stderr)
    from collections import Counter
    print("survivors by vendor:", dict(Counter(r["vendor"] for r in surv)), file=sys.stderr)

if __name__ == "__main__":
    main()