`visual-search/pattern_siblings.py`

131 lines
#!/usr/bin/env python3
"""Pattern-level sibling stage for visual-search (TK-12090 Lane P) — OPTIONAL, default OFF.

When a photo matches one COLORWAY, also return the other colorways of the same PATTERN.
Ported (not forked) from ~/Projects/dw-pairs-well/tools:
  * signature  = ci-hash-build.py        dHash(8) of grayscale + dHash of the tonal inverse
  * distance   = pattern-id-titlefirst-final.py ci_dist: min of the 4 cross Hamming distances,
                 so light-on-dark and dark-on-light colorways of one motif still match
  * threshold  = HAM_THRESH default 30 (WS-1 value chosen 2026-06-26 in that script)
  * grouping   = the WS-1 title-first + CI-confirmed map pattern-id-titlefirst-by-dwsku.json

Candidate siblings come from (1) the WS-1 pattern_id group, else (2) same vendor + same normalized
pattern_name in the search index. Each candidate is then CONFIRMED by the colour-invariant distance
when both hashes are known; unhashed candidates are returned as confirmed=None (never guessed).

Env: PAIRS_WELL_TOOLS (default ~/Projects/dw-pairs-well/tools), HAM_THRESH (30),
     PATTERN_SIBLINGS_FETCH=1 to hash missing images on the fly (bounded, disk-cached, $0).
"""
import hashlib, io, json, os, re, ssl, urllib.request
from collections import defaultdict
from PIL import Image, ImageOps

TOOLS = os.environ.get("PAIRS_WELL_TOOLS", os.path.expanduser("~/Projects/dw-pairs-well/tools"))
THRESH = int(os.environ.get("HAM_THRESH", "30"))
FETCH = os.environ.get("PATTERN_SIBLINGS_FETCH") == "1"
CACHE = os.path.expanduser("~/.cache/ci-hash")          # same cache dir as ci-hash-build.py
HASH = 8

# ── signature + distance: verbatim logic from dw-pairs-well/tools ─────────────────────────
def dhash(img, n=HASH):
    g = img.convert("L").resize((n + 1, n), Image.LANCZOS)
    px = list(g.getdata()); v = 0
    for r in range(n):
        b = r * (n + 1)
        for c in range(n):
            v = (v << 1) | (1 if px[b + c] > px[b + c + 1] else 0)
    return f"{v:016x}"

def inv_dhash(img, n=HASH):
    return dhash(ImageOps.invert(img.convert("L")), n)

def ham(a, b): return bin(int(a, 16) ^ int(b, 16)).count("1")

def ci_dist(A, B):
    return min(ham(A[0], B[0]), ham(A[1], B[1]), ham(A[0], B[1]), ham(A[1], B[0]))

# ── data ─────────────────────────────────────────────────────────────────────────────────
_D = {"loaded": False}

def _norm(s): return re.sub(r"[^a-z0-9]+", " ", (s or "").lower()).strip()

def load():
    if _D["loaded"]: return _D
    pid, members, url_of, ci = {}, defaultdict(list), {}, {}
    try:
        for k, v in json.load(open(os.path.join(TOOLS, "pattern-id-titlefirst-by-dwsku.json"))).items():
            if v and v.get("pid"): pid[k] = v["pid"]; members[v["pid"]].append(k)
    except OSError: pass
    try:
        for line in open(os.path.join(TOOLS, "_active-for-ci.tsv")):
            p = line.rstrip("\n").split("\t")
            if len(p) >= 5 and p[4].startswith("http"): url_of.setdefault(p[0], p[4])
    except OSError: pass
    try:
        for line in open(os.path.join(TOOLS, "ci-hashes.jsonl")):
            try:
                r = json.loads(line)
                if r.get("status") == "ok": ci[r["url"]] = (r["dhash"], r["inv"])
            except ValueError: pass
    except OSError: pass
    _D.update(loaded=True, pid=pid, members=members, url_of=url_of, ci=ci)
    return _D

def _fetch_sig(url):
    os.makedirs(CACHE, exist_ok=True)
    key = os.path.join(CACHE, hashlib.md5(url.encode()).hexdigest() + ".bin")
    if not (os.path.exists(key) and os.path.getsize(key) > 0):
        ctx = ssl.create_default_context()
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 dw-ci-hash"})
        with urllib.request.urlopen(req, timeout=15, context=ctx) as r: open(key, "wb").write(r.read())
    img = Image.open(io.BytesIO(open(key, "rb").read())); img.load()
    return (dhash(img), inv_dhash(img))

def signature(dw_sku, image_url=None, budget=None):
    d = load()
    url = d["url_of"].get(dw_sku) or image_url
    if not url: return None
    if url in d["ci"]: return d["ci"][url]
    if FETCH and budget is not None and budget[0] > 0:
        budget[0] -= 1
        try:
            sig = _fetch_sig(url); d["ci"][url] = sig; return sig
        except Exception:
            return None
    return None

def siblings(match, index_meta=None, limit=24):
    """match: a /search result dict ({dw_sku, vendor, pattern, image}). index_meta: the service's
    STATE['meta'] list, used for the vendor+pattern_name fallback. Returns a dict for the response."""
    d = load(); sku = match.get("dw_sku") or ""
    group_key, cands, via = None, [], None
    if sku in d["pid"]:
        group_key = d["pid"][sku]; via = "ws1_pattern_id"
        cands = [(s, None) for s in d["members"][group_key] if s != sku]
    elif index_meta and match.get("pattern"):
        pn, vn = _norm(match.get("pattern")), match.get("vendor")
        group_key = f"{vn}::{pn}"; via = "vendor+pattern_name"
        seen = {sku}
        for m in index_meta:
            s = m.get("dw_sku")
            if s and s not in seen and m.get("vendor") == vn and _norm(m.get("pattern")) == pn:
                seen.add(s); cands.append((s, m.get("image")))
    budget = [limit]
    base = signature(sku, match.get("image"), budget)
    out = []
    for s, img in cands[: limit * 4]:
        sig = signature(s, img, budget)
        dist = ci_dist(base, sig) if (base and sig) else None
        out.append({"dw_sku": s, "ci_dist": dist, "confirmed": (dist <= THRESH) if dist is not None else None})
    out.sort(key=lambda x: (x["confirmed"] is not True, x["ci_dist"] if x["ci_dist"] is not None else 99))
    reason = None if via else "no WS-1 pattern_id and no pattern_name on the match — siblings NOT MEASURED"
    return {"group": group_key, "via": via, "reason": reason, "thresh": THRESH, "query_hashed": base is not None,
            "n": len(out), "confirmed": sum(1 for x in out if x["confirmed"]),
            "rejected": sum(1 for x in out if x["confirmed"] is False), "siblings": out[:limit]}

if __name__ == "__main__":
    import sys
    for s in sys.argv[1:]:
        r = siblings({"dw_sku": s})
        print(json.dumps({k: v for k, v in r.items() if k != "siblings"}), [ (x["dw_sku"], x["ci_dist"]) for x in r["siblings"][:8]])