← back to Dw Photo Capture
visual-search/search_service.py
283 lines
#!/usr/bin/env python3
# Visual product search service. Holds the catalog CLIP embeddings in RAM and answers nearest-image
# queries. dwphoto (Node) POSTs a photo → this returns the visually-closest catalog products.
# No pgvector needed: brute-force cosine over the normalized matrix (numpy, ~ms for hundreds of k).
#
# DW_UNIFIED_DB=<uri> VS_PORT=9914 python3 search_service.py
import os, io, json, base64, threading, time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import numpy as np, psycopg2
from PIL import Image
# torch/open_clip are imported LAZILY (see _get_model) — only the /search path re-embeds a photo
# and needs the model. /similar, /health and /reload read STORED vectors only, so the service can
# run (and /similar can be proved) on a box without torch installed.
DB = os.environ["DW_UNIFIED_DB"]
PORT = int(os.environ.get("VS_PORT", "9914"))
# Bind host. Defaults to 127.0.0.1 (localhost-only, unchanged behavior). Set
# VS_HOST=0.0.0.0 to expose on the tailnet so a remote host (e.g. Kamatera
# pairs-well) can reach it — only do that when the exposure is intended.
HOST = os.environ.get("VS_HOST", "127.0.0.1")
_MODEL = {"model": None, "preprocess": None}
def _get_model():
"""Load CLIP on first /search only. Uses the SAME weights (ViT-B-32/laion2b + optional
dw_clip_ft.pt) the catalog was embedded with, so re-embedded photos share the catalog's space."""
if _MODEL["model"] is None:
import torch, open_clip
print("loading CLIP…", flush=True)
model, _, preprocess = open_clip.create_model_and_transforms("ViT-B-32", pretrained="laion2b_s34b_b79k")
_ft = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dw_clip_ft.pt")
if os.path.exists(_ft):
model.load_state_dict(torch.load(_ft, map_location="cpu"), strict=False)
print("loaded fine-tuned DW CLIP weights:", _ft, flush=True)
model.eval(); torch.set_num_threads(max(1, os.cpu_count() - 2))
_MODEL["model"], _MODEL["preprocess"] = model, preprocess
return _MODEL["model"], _MODEL["preprocess"]
LOCK = threading.Lock()
# SKU_IX maps dw_sku -> row index in M so /similar can look up a product's OWN stored vector
# with NO re-embed (the vector was already written by embed_catalog.py using the SAME weights).
STATE = {"M": None, "meta": [], "loaded_at": 0, "n": 0, "sku_ix": {}, "fam": [], "styles": []}
# ── HEX → color-family bucket (server-side port of Fork A's familyFromHex, commit 7cf6b2b,
# snippets/color-palette.liquid). SAME HSL thresholds so the color-dot the shopper clicks
# and the /similar re-rank agree on what "family" a hex is. Do NOT drift these. ──
def _family_from_hex(hex_str):
if not hex_str: return None
h = hex_str.strip().lstrip("#")
if len(h) == 3: h = "".join(c*2 for c in h)
if len(h) < 6: return None
try:
r, g, b = int(h[0:2],16)/255.0, int(h[2:4],16)/255.0, int(h[4:6],16)/255.0
except ValueError:
return None
mx, mn = max(r,g,b), min(r,g,b); d = mx-mn
l = (mx+mn)/2.0
s = 0.0 if d == 0 else d/(1-abs(2*l-1)) if (1-abs(2*l-1)) else 0.0
hd = 0.0
if d != 0:
if mx == r: hd = 60*(((g-b)/d) % 6)
elif mx == g: hd = 60*(((b-r)/d) + 2)
else: hd = 60*(((r-g)/d) + 4)
if hd < 0: hd += 360
if s < 0.18:
if l > 0.86: return "white"
if l < 0.18: return "black"
if 20 <= hd <= 70 and l > 0.30: return "beige"
return "gray"
if l > 0.80 and 20 <= hd < 70: return "beige"
if l < 0.40 and s < 0.75 and ((10 <= hd < 70) or hd >= 345): return "brown"
if l < 0.16: return "black"
if hd < 15 or hd >= 345: return "red"
if hd < 40: return "beige" if (s < 0.45 and l > 0.55) else "orange"
if hd < 55: return "beige" if s < 0.45 else ("gold" if hd < 50 else "yellow")
if hd < 65: return "gold" if l < 0.45 else "yellow"
if hd < 170: return "green"
if hd < 200: return "teal"
if hd < 255: return "blue"
if hd < 290: return "purple"
return "pink"
# families that read as "close enough" so the soft-boost doesn't punish a near-neighbor
_FAM_ADJ = {
"beige": {"brown", "gold", "white", "gray"}, "brown": {"beige", "gold", "orange"},
"gold": {"beige", "yellow", "brown", "orange"}, "yellow": {"gold", "orange"},
"orange": {"brown", "gold", "red"}, "red": {"orange", "pink"},
"pink": {"red", "purple"}, "purple": {"pink", "blue"}, "blue": {"purple", "teal"},
"teal": {"blue", "green"}, "green": {"teal"}, "gray": {"white", "black", "beige"},
"white": {"beige", "gray"}, "black": {"gray"},
}
def load():
conn = psycopg2.connect(DB); cur = conn.cursor()
# LEFT JOIN vendor_catalog for the color/style re-rank signals (dominant_color_hex + ai_styles),
# and shopify_products for the storefront `handle` so /similar can return a real PDP link.
# Keyed by dw_sku; rows without a match just carry null fam/styles/handle (soft signals, never required).
# handle joins ONLY to ACTIVE + Online-Store-published products: /similar's DW-safe
# guard drops handleless rows, so archived/draft/unpublished products can never be
# suggested as clickable results (their vectors still serve as query anchors).
cur.execute("""select e.vc_id, e.dw_sku, e.mfr_sku, e.vendor_code, e.pattern_name, e.image_url,
e.embedding, v.dominant_color_hex, v.ai_styles, sp.handle
from image_embeddings e
left join vendor_catalog v on v.dw_sku = e.dw_sku
left join shopify_products sp on sp.dw_sku = e.dw_sku
and sp.status = 'ACTIVE' and sp.online_store_published is true
where e.embedding is not null""")
vecs, meta, fam, styles, sku_ix = [], [], [], [], {}
for vc_id, dw_sku, mfr, vc, pat, url, emb, hexc, ai_styles, handle in cur:
i = len(meta)
vecs.append(np.frombuffer(bytes(emb), dtype="float32"))
meta.append({"vc_id": vc_id, "dw_sku": dw_sku, "mfr_sku": mfr, "vendor": vc, "pattern": pat, "image": url, "handle": handle})
fam.append(_family_from_hex(hexc))
# ai_styles is jsonb array → normalize to a lowercase set for overlap test
sset = set()
if ai_styles:
try:
arr = ai_styles if isinstance(ai_styles, list) else json.loads(ai_styles)
sset = {str(x).strip().lower() for x in arr if x}
except Exception:
pass
styles.append(sset)
if dw_sku and dw_sku not in sku_ix: # first vector wins for a dw_sku
sku_ix[dw_sku] = i
cur.close(); conn.close()
M = np.vstack(vecs).astype("float32") if vecs else np.zeros((0, 512), "float32")
with LOCK:
STATE["M"] = M; STATE["meta"] = meta; STATE["loaded_at"] = time.time(); STATE["n"] = len(meta)
STATE["sku_ix"] = sku_ix; STATE["fam"] = fam; STATE["styles"] = styles
print(f"loaded {len(meta)} embeddings · {len(sku_ix)} skus indexed", flush=True)
def embed_photo(b64):
import torch
model, preprocess = _get_model()
img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
with torch.no_grad():
v = model.encode_image(preprocess(img).unsqueeze(0))
v = (v / v.norm(dim=-1, keepdim=True)).cpu().numpy().astype("float32")[0]
return v
class H(BaseHTTPRequestHandler):
def log_message(self, *a): pass
def _send(self, code, obj):
b = json.dumps(obj).encode(); self.send_response(code)
self.send_header("Content-Type", "application/json"); self.send_header("Content-Length", str(len(b)))
self.end_headers(); self.wfile.write(b)
def do_GET(self):
if self.path.startswith("/health"): return self._send(200, {"ok": True, "n": STATE["n"], "loaded_at": STATE["loaded_at"]})
if self.path.startswith("/reload"):
try: load(); return self._send(200, {"ok": True, "n": STATE["n"]})
except Exception as e: return self._send(500, {"ok": False, "err": str(e)})
self._send(404, {"err": "not found"})
def do_POST(self):
if self.path.startswith("/similar"): return self._similar()
if self.path.startswith("/pairwise"): return self._pairwise()
if not self.path.startswith("/search"): return self._send(404, {"err": "not found"})
try:
n = int(self.headers.get("Content-Length", 0)); p = json.loads(self.rfile.read(n) or b"{}")
b64 = (p.get("image") or "").split(",")[-1]
k = int(p.get("k", 8))
if not b64: return self._send(400, {"err": "image required"})
q = embed_photo(b64)
with LOCK: M, meta = STATE["M"], STATE["meta"]
if M is None or not len(M): return self._send(200, {"ok": True, "results": [], "n": 0})
sims = M @ q # cosine (both L2-normalized)
idx = np.argpartition(-sims, min(k, len(sims)-1))[:k]
idx = idx[np.argsort(-sims[idx])]
res = [{**meta[i], "score": round(float(sims[i]), 4)} for i in idx]
self._send(200, {"ok": True, "n": STATE["n"], "results": res})
except Exception as e:
self._send(500, {"ok": False, "err": str(e)})
def _pairwise(self):
# Bulk cosine: source dw_sku vs an explicit candidate list. Contract:
# {dw_sku, candidates: [dw_sku, …]} → {ok, cosines: {sku: cos}}.
# Consumer: dw-pairs-well /api/pairs blends these as a bounded style-world
# sub-score (mid-similarity sweet spot; near-duplicates get penalized there).
# Candidates without a stored vector are simply absent from the map (neutral).
try:
n = int(self.headers.get("Content-Length", 0)); p = json.loads(self.rfile.read(n) or b"{}")
dw_sku = (p.get("dw_sku") or "").strip()
cands = p.get("candidates") or []
if not dw_sku or not isinstance(cands, list):
return self._send(400, {"ok": False, "err": "dw_sku + candidates[] required"})
cands = [str(c).strip() for c in cands[:5000] if c]
with LOCK:
M, sku_ix = STATE["M"], STATE["sku_ix"]
if M is None or not len(M):
return self._send(200, {"ok": True, "cosines": {}})
qi = sku_ix.get(dw_sku)
if qi is None:
return self._send(404, {"ok": False, "err": f"no stored embedding for dw_sku {dw_sku}"})
q = M[qi]
ixs, skus = [], []
for c in cands:
ci = sku_ix.get(c)
if ci is not None and ci != qi:
ixs.append(ci); skus.append(c)
if not ixs:
return self._send(200, {"ok": True, "cosines": {}})
sims = M[np.array(ixs)] @ q
out = {s: round(float(v), 4) for s, v in zip(skus, sims)}
self._send(200, {"ok": True, "n": len(out), "cosines": out})
except Exception as e:
self._send(500, {"ok": False, "err": str(e)})
def _similar(self):
# PDP color-dot → CLIP visual similarity. Contract: {dw_sku, hex, style, k}.
# Looks up the product's OWN stored embedding by dw_sku (NO re-embed — the vector was
# written by embed_catalog.py with the SAME ViT-B-32/laion2b (+ optional dw_clip_ft.pt)
# weights this process loads, so query and catalog live in the identical vector space).
# Then cosine-ranks a generous pool and SOFT re-ranks by color-family + style overlap so
# strict filters never empty the result set.
try:
n = int(self.headers.get("Content-Length", 0)); p = json.loads(self.rfile.read(n) or b"{}")
dw_sku = (p.get("dw_sku") or "").strip()
hexq = (p.get("hex") or "").strip()
style_q = p.get("style") or ""
k = max(1, min(int(p.get("k", 8)), 50))
if not dw_sku:
return self._send(400, {"ok": False, "err": "dw_sku required"})
with LOCK:
M, meta = STATE["M"], STATE["meta"]
sku_ix, fam, styles = STATE["sku_ix"], STATE["fam"], STATE["styles"]
if M is None or not len(M):
return self._send(200, {"ok": True, "n": 0, "results": []})
qi = sku_ix.get(dw_sku)
if qi is None:
# clear error, never crash — the sku has no stored vector (never embedded / image failed)
return self._send(404, {"ok": False, "err": f"no stored embedding for dw_sku {dw_sku}"})
q = M[qi] # the product's OWN vector (already unit-norm)
sims = M @ q # cosine over the full matrix
# generous candidate pool (10x k, capped) BEFORE soft filters, so re-rank has room
pool = max(k * 10, 200)
cand = np.argpartition(-sims, min(pool, len(sims)-1))[:pool]
# normalize the query's color-family + style set for the soft signals
qfam = _family_from_hex(hexq) if hexq else (fam[qi] if qi < len(fam) else None)
qstyles = set()
if isinstance(style_q, list):
qstyles = {str(x).strip().lower() for x in style_q if x}
elif isinstance(style_q, str) and style_q.strip():
qstyles = {s.strip().lower() for s in style_q.replace(",", " ").split() if s.strip()}
if not qstyles and qi < len(styles):
qstyles = styles[qi]
scored = []
for i in cand:
if i == qi: continue # exclude the query itself
base = float(sims[i]) # visual cosine is the dominant term
boost = 0.0
cfam = fam[i] if i < len(fam) else None
if qfam and cfam:
if cfam == qfam: boost += 0.06 # exact family match
elif cfam in _FAM_ADJ.get(qfam, ()): boost += 0.03 # adjacent family (near hue)
else: boost -= 0.03 # off-family, gentle nudge down
if qstyles and i < len(styles) and styles[i]:
if qstyles & styles[i]: boost += 0.04 # any shared style term
scored.append((base + boost, base, i))
scored.sort(key=lambda t: t[0], reverse=True)
res = []
for _rank, base, i in scored: # iterate ALL (sorted desc), backfill to k good cards
if len(res) >= k: break
m = meta[i]
h = m.get("handle"); img = m.get("image") or ""
# DW-safe guard (customer-facing): only real storefront products with a DW-hosted image.
# Drops (1) handleless index entries (not real PDPs → broken /search-fallback cards) and
# (2) raw-vendor hotlinks like www.wallquest.com — those 404 (broken images) AND leak the
# private-label vendor to customers (WallQuest→Malibu). Backfill keeps the strip full.
if not h: continue
if "cdn.shopify.com" not in img and "designerwallcoverings.com" not in img: continue
res.append({"dw_sku": m["dw_sku"], "image": img, "pattern": m["pattern"],
"vendor": m["vendor"], "handle": h, "score": round(base, 4)})
self._send(200, {"ok": True, "n": STATE["n"], "query": {
"dw_sku": dw_sku, "family": qfam, "styles": sorted(qstyles)}, "results": res})
except Exception as e:
self._send(500, {"ok": False, "err": str(e)})
if __name__ == "__main__":
load()
print(f"visual-search on :{PORT} · {STATE['n']} embeddings", flush=True)
ThreadingHTTPServer((HOST, PORT), H).serve_forever()