260 lines
11 KiB
Python
260 lines
11 KiB
Python
"""
|
|
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 (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
|
|
|
|
# Fixed seed for both the pixel sample and k-means' own center
|
|
# initialization (cv2.KMEANS_PP_CENTERS is randomized). Without this, the
|
|
# *same* photo could come back with slightly different dominant colors on
|
|
# different requests/server restarts, which made the color match silently
|
|
# non-reproducible -- not the actual bug behind bad matches, but worth
|
|
# closing since it made every other bug here harder to pin down.
|
|
_KMEANS_SEED = 20260101
|
|
_rng = np.random.default_rng(_KMEANS_SEED)
|
|
|
|
# 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 -> {"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):
|
|
"""K-means dominant-color extraction, clustered in Lab space (not raw
|
|
BGR). BGR/RGB Euclidean distance conflates brightness with hue, so
|
|
clustering there tends to waste cluster budget separating a highlight
|
|
from a shadow on the *same* material (e.g. two "white ribbon" clusters)
|
|
instead of separating genuinely different colors (e.g. the ribbon from
|
|
the roses) -- exactly the failure mode that was drowning out small but
|
|
decisive colors. Lab is perceptually uniform, so distance in Lab space
|
|
lines up with what "looks different" to begin with, and it's also the
|
|
space every downstream comparison already measures in.
|
|
|
|
Returns each color's own Lab center (used for matching) alongside its
|
|
RGB conversion (display-only)."""
|
|
pixels = _masked_pixels_bgr(bgr, mask)
|
|
if len(pixels) == 0:
|
|
return []
|
|
|
|
if len(pixels) > MAX_KMEANS_SAMPLE:
|
|
idx = _rng.choice(len(pixels), MAX_KMEANS_SAMPLE, replace=False)
|
|
pixels = pixels[idx]
|
|
|
|
# OpenCV's uint8 BGR2LAB packs L into 0-255 and offsets a/b by +128 (a
|
|
# storage convenience for 8-bit images), *not* the standard CIE Lab
|
|
# scale (L: 0-100, a/b roughly -128..127) that MAX_LAB_DISTANCE below
|
|
# assumes. Left unconverted, that mismatch means "maximally different"
|
|
# (e.g. black vs. white) measures 255, not ~100 -- so real distances
|
|
# blow past the similarity floor constantly and every comparison
|
|
# collapses toward 0%, destroying exactly the separation this is
|
|
# supposed to produce. Converting here keeps every distance computed
|
|
# against these centers (k-means itself, and every comparison below)
|
|
# correctly calibrated.
|
|
lab_pixels = cv2.cvtColor(pixels.reshape(-1, 1, 3), cv2.COLOR_BGR2LAB).reshape(-1, 3).astype(np.float32)
|
|
lab_pixels[:, 0] *= 100.0 / 255.0
|
|
lab_pixels[:, 1] -= 128.0
|
|
lab_pixels[:, 2] -= 128.0
|
|
|
|
unique_count = len(np.unique(lab_pixels, axis=0))
|
|
k_eff = min(k, unique_count)
|
|
if k_eff < 1:
|
|
return []
|
|
|
|
cv2.setRNGSeed(_KMEANS_SEED)
|
|
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 0.5)
|
|
_, labels, centers = cv2.kmeans(lab_pixels, 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):
|
|
lab = centers[i]
|
|
# invert back to OpenCV's 0-255 uint8 convention just for display
|
|
l_cv = np.clip(round(lab[0] * 255.0 / 100.0), 0, 255)
|
|
a_cv = np.clip(round(lab[1] + 128.0), 0, 255)
|
|
b_cv = np.clip(round(lab[2] + 128.0), 0, 255)
|
|
lab_u8 = np.array([l_cv, a_cv, b_cv], dtype=np.uint8)
|
|
bgr_center = cv2.cvtColor(lab_u8.reshape(1, 1, 3), cv2.COLOR_LAB2BGR)[0][0]
|
|
b, g, r = bgr_center
|
|
colors.append({
|
|
"rgb": [int(r), int(g), int(b)],
|
|
"lab": [float(lab[0]), float(lab[1]), float(lab[2])],
|
|
"pct": round(float(counts[i]) / total * 100, 1),
|
|
})
|
|
colors.sort(key=lambda c: c["pct"], reverse=True)
|
|
return colors
|
|
|
|
|
|
def _lab_distance(lab_a, lab_b):
|
|
return float(np.linalg.norm(np.asarray(lab_a, dtype=np.float32) -
|
|
np.asarray(lab_b, dtype=np.float32)))
|
|
|
|
|
|
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 %.
|
|
Purely for the human-readable "which color matched which" breakdown
|
|
shown in the UI -- greedy nearest-neighbor can let several input colors
|
|
point at the same template color while leaving others unmatched, which
|
|
is fine for an explanation but not for scoring (see
|
|
compare_input_to_templates)."""
|
|
pairs = []
|
|
for ic in input_colors:
|
|
if not template_colors:
|
|
break
|
|
best = min(template_colors, key=lambda tc: _lab_distance(ic["lab"], tc["lab"]))
|
|
dist = _lab_distance(ic["lab"], best["lab"])
|
|
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 _emd_similarity(input_colors, template_colors):
|
|
"""Overall match %, via Earth Mover's Distance between the two color
|
|
sets treated as weighted signatures (weight = pct of photo, position =
|
|
Lab coordinate).
|
|
|
|
This replaces a greedy "each input color grabs its single nearest
|
|
template color, then average" scheme, which has two real problems:
|
|
it's a many-to-one assignment (several input colors can all latch onto
|
|
the same template color while other template colors -- e.g. a second
|
|
flower type -- never factor in at all), and it only asks "does the
|
|
template have something close to each of my colors", never the reverse.
|
|
EMD instead finds the cheapest way to reshape the whole input
|
|
distribution into the whole template distribution, using every cluster
|
|
on both sides, weighted by how much of each photo it actually is."""
|
|
if not input_colors or not template_colors:
|
|
return 0.0
|
|
sig_a = np.array([[c["pct"]] + c["lab"] for c in input_colors], dtype=np.float32)
|
|
sig_b = np.array([[c["pct"]] + c["lab"] for c in template_colors], dtype=np.float32)
|
|
dist, _, _ = cv2.EMD(sig_a, sig_b, cv2.DIST_L2)
|
|
similarity = max(0.0, 100.0 * (1 - dist / MAX_LAB_DISTANCE))
|
|
return round(similarity, 1)
|
|
|
|
|
|
def set_template_color_data(name, bgr, mask):
|
|
# name -> [{"hist","colors"}, ...] -- index 0 is always the canonical
|
|
# template; further entries are hidden "hard example" photos (see
|
|
# config.HARD_EXAMPLES_DIR) matched too but never shown in the UI.
|
|
_template_color_data[name] = [{
|
|
"hist": compute_hs_histogram(bgr, mask),
|
|
"colors": dominant_colors(bgr, mask),
|
|
}]
|
|
|
|
|
|
def add_hard_example_color_data(name, bgr, mask):
|
|
_template_color_data[name].append({
|
|
"hist": compute_hs_histogram(bgr, mask),
|
|
"colors": dominant_colors(bgr, mask),
|
|
})
|
|
|
|
|
|
def _blended_match_pct(input_hist, input_colors, data):
|
|
hist_pct = histogram_match_pct(input_hist, data["hist"])
|
|
emd_pct = _emd_similarity(input_colors, data["colors"])
|
|
# The full HS histogram is empirically the stronger signal on its own
|
|
# (it sees the whole color distribution, not just 5 clusters), but it
|
|
# has one specific blind spot: two photos that both happen to be
|
|
# dominated by white/cream (ribbon, pale petals) can score a high
|
|
# overlap there even when their *actual* flower colors differ
|
|
# completely, since the histogram has no notion of "this hue only
|
|
# covers a little of the photo, that one covers a lot". The dominant-
|
|
# color EMD score is weaker on average but is built exactly to catch
|
|
# that case (it's explicitly weighted by how much of the photo each
|
|
# color is). Averaging keeps the histogram's overall strength while
|
|
# letting EMD pull the score down when a would-be match is really just
|
|
# "both photos are mostly white".
|
|
return round((hist_pct + emd_pct) / 2, 1), hist_pct, emd_pct
|
|
|
|
|
|
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, gallery in _template_color_data.items():
|
|
# Each template is matched against its whole gallery (canonical
|
|
# photo + any hidden hard examples); the best-scoring entry wins --
|
|
# a real photo that looks more like one of the hard examples than
|
|
# the canonical shot still gets full credit, and the color-pairs
|
|
# breakdown shown in the UI comes from that same winning entry so
|
|
# the numbers and the explanation stay consistent.
|
|
best = None
|
|
for data in gallery:
|
|
match_pct, hist_pct, emd_pct = _blended_match_pct(input_hist, input_colors, data)
|
|
if best is None or match_pct > best[0]:
|
|
best = (match_pct, hist_pct, emd_pct, data)
|
|
match_pct, hist_pct, emd_pct, winning_data = best
|
|
pairs = match_color_pairs(input_colors, winning_data["colors"])
|
|
results.append({
|
|
"template": name,
|
|
"match_pct": match_pct,
|
|
"hist_match_pct": hist_pct,
|
|
"dominant_color_match_pct": emd_pct,
|
|
"color_pairs": pairs,
|
|
})
|
|
results.sort(key=lambda r: r["match_pct"], reverse=True)
|
|
|
|
return {
|
|
"input_dominant_colors": input_colors,
|
|
"templates": results,
|
|
}
|