149 lines
5.2 KiB
Python
149 lines
5.2 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 (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,
|
|
}
|