Add project files

pipeline/color.py

@@ -0,0 +1,148 @@
+"""
+Independent color-space comparison: how much of the uploaded photo's color
+distribution is also present in each template, plus which specific colors
+matched. Entirely separate from the SIFT/ORB/SuperGlue/LoFTR scores and the
+weighted final match -- nothing here feeds into those, it's its own section.
+
+Template-side color data (histogram + dominant colors) is precomputed once
+at startup from the same masked bgr/mask already produced during bootstrap,
+same principle as the other methods: a request only ever analyzes the one
+uploaded image.
+"""
+
+import cv2
+import numpy as np
+
+N_DOMINANT_COLORS = 5
+HIST_HUE_BINS = 30
+HIST_SAT_BINS = 32
+
+# k-means on every foreground pixel is wasteful and, given this app's
+# history of memory blowups, worth bounding explicitly -- a random sample
+# is statistically equivalent for "what are the dominant colors" purposes.
+MAX_KMEANS_SAMPLE = 20000
+
+# Empirical ceiling for CIE76 Lab distance beyond which two colors are
+# considered maximally dissimilar (similarity floors at 0%). ~100 comfortably
+# covers the largest perceptual differences (e.g. black vs. white is ~100).
+MAX_LAB_DISTANCE = 100.0
+
+_template_color_data = {}  # name -> {"hist": ndarray, "colors": [...]}
+
+
+def _masked_pixels_bgr(bgr, mask):
+    ys, xs = np.where(mask > 0)
+    if len(ys) == 0:
+        return np.zeros((0, 3), dtype=np.uint8)
+    return bgr[ys, xs]
+
+
+def compute_hs_histogram(bgr, mask):
+    """Hue+Saturation 2D histogram (Value/brightness deliberately excluded
+    so lighting differences between photos don't masquerade as color
+    differences), normalized so it sums to 1."""
+    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
+    hist = cv2.calcHist([hsv], [0, 1], mask, [HIST_HUE_BINS, HIST_SAT_BINS],
+                        [0, 180, 0, 256])
+    total = hist.sum()
+    if total > 0:
+        hist = hist / total
+    return hist
+
+
+def histogram_match_pct(hist_a, hist_b):
+    """% of one distribution's mass that the other also covers -- both
+    histograms are pre-normalized, so HISTCMP_INTERSECT directly sums
+    min(a[i], b[i]) into a 0-1 overlap fraction."""
+    intersection = cv2.compareHist(hist_a.astype(np.float32), hist_b.astype(np.float32),
+                                    cv2.HISTCMP_INTERSECT)
+    return round(float(intersection) * 100, 1)
+
+
+def dominant_colors(bgr, mask, k=N_DOMINANT_COLORS):
+    pixels = _masked_pixels_bgr(bgr, mask)
+    if len(pixels) == 0:
+        return []
+
+    if len(pixels) > MAX_KMEANS_SAMPLE:
+        idx = np.random.choice(len(pixels), MAX_KMEANS_SAMPLE, replace=False)
+        pixels = pixels[idx]
+
+    unique_count = len(np.unique(pixels, axis=0))
+    k_eff = min(k, unique_count)
+    if k_eff < 1:
+        return []
+
+    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 0.5)
+    _, labels, centers = cv2.kmeans(pixels.astype(np.float32), k_eff, None,
+                                     criteria, 3, cv2.KMEANS_PP_CENTERS)
+
+    labels = labels.flatten()
+    counts = np.bincount(labels, minlength=k_eff)
+    total = counts.sum()
+
+    colors = []
+    for i in range(k_eff):
+        b, g, r = centers[i]
+        colors.append({
+            "rgb": [int(round(r)), int(round(g)), int(round(b))],
+            "pct": round(float(counts[i]) / total * 100, 1),
+        })
+    colors.sort(key=lambda c: c["pct"], reverse=True)
+    return colors
+
+
+def _lab_distance(rgb_a, rgb_b):
+    a = np.uint8([[[rgb_a[2], rgb_a[1], rgb_a[0]]]])  # rgb -> bgr for cv2
+    b = np.uint8([[[rgb_b[2], rgb_b[1], rgb_b[0]]]])
+    lab_a = cv2.cvtColor(a, cv2.COLOR_BGR2LAB)[0][0].astype(np.float32)
+    lab_b = cv2.cvtColor(b, cv2.COLOR_BGR2LAB)[0][0].astype(np.float32)
+    return float(np.linalg.norm(lab_a - lab_b))
+
+
+def match_color_pairs(input_colors, template_colors):
+    """For each of the input's dominant colors, find the closest template
+    dominant color by perceptual (Lab) distance and report a similarity %."""
+    pairs = []
+    for ic in input_colors:
+        if not template_colors:
+            break
+        best = min(template_colors, key=lambda tc: _lab_distance(ic["rgb"], tc["rgb"]))
+        dist = _lab_distance(ic["rgb"], best["rgb"])
+        similarity = max(0.0, 100.0 * (1 - dist / MAX_LAB_DISTANCE))
+        pairs.append({
+            "input_rgb": ic["rgb"],
+            "input_pct": ic["pct"],
+            "template_rgb": best["rgb"],
+            "template_pct": best["pct"],
+            "similarity": round(similarity, 1),
+        })
+    return pairs
+
+
+def set_template_color_data(name, bgr, mask):
+    _template_color_data[name] = {
+        "hist": compute_hs_histogram(bgr, mask),
+        "colors": dominant_colors(bgr, mask),
+    }
+
+
+def compare_input_to_templates(input_bgr, input_mask):
+    input_hist = compute_hs_histogram(input_bgr, input_mask)
+    input_colors = dominant_colors(input_bgr, input_mask)
+
+    results = []
+    for name, data in _template_color_data.items():
+        match_pct = histogram_match_pct(input_hist, data["hist"])
+        pairs = match_color_pairs(input_colors, data["colors"])
+        results.append({
+            "template": name,
+            "match_pct": match_pct,
+            "color_pairs": pairs,
+        })
+    results.sort(key=lambda r: r["match_pct"], reverse=True)
+
+    return {
+        "input_dominant_colors": input_colors,
+        "templates": results,
+    }