Upload 2 files

color_calibration.py

@@ -0,0 +1,140 @@
+"""
+Color calibration via X-Rite ColorChecker Classic (24 patches).
+
+Detects the checker in both a reference and target image, computes a
+white-balance + CCM + luminance correction pipeline, and applies it to
+the full target image.
+"""
+
+import cv2 as cv
+import numpy as np
+
+
+# ── linear / sRGB helpers ──────────────────────────────────────────
+
+def to_linear_srgb(u8_bgr):
+    rgb = cv.cvtColor(u8_bgr, cv.COLOR_BGR2RGB).astype(np.float32) / 255.0
+    a = 0.055
+    lin = np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + a) / (1 + a)) ** 2.4)
+    return lin
+
+
+def to_srgb_u8(lin_rgb):
+    a = 0.055
+    srgb = np.where(lin_rgb <= 0.0031308, 12.92 * lin_rgb, (1 + a) * np.power(lin_rgb, 1/2.4) - a)
+    srgb = np.clip(srgb, 0, 1)
+    return cv.cvtColor((srgb * 255.0).astype(np.uint8), cv.COLOR_RGB2BGR)
+
+
+# ── checker geometry ───────────────────────────────────────────────
+
+def warp_checker(img, corners, out_w=600, out_h=400):
+    dst = np.float32([[0, 0], [out_w-1, 0], [out_w-1, out_h-1], [0, out_h-1]])
+    H = cv.getPerspectiveTransform(np.float32(corners), dst)
+    warped = cv.warpPerspective(img, H, (out_w, out_h), flags=cv.INTER_CUBIC)
+    return warped
+
+
+def detect_checker_corners(img_bgr):
+    det = cv.mcc.CCheckerDetector_create()
+    ok = det.process(img_bgr, cv.mcc.MCC24)
+    if not ok:
+        raise RuntimeError("ColorChecker not found")
+    lst = det.getListColorChecker()
+    cc = lst[0]
+    if hasattr(cc, "getBox"):
+        corners = np.array(cc.getBox(), dtype=np.float32)
+    else:
+        corners = np.array(cc.getCorners(), dtype=np.float32)
+    return corners
+
+
+def sample_24_patches(warped, margin=12):
+    H, W = warped.shape[:2]
+    cell_w, cell_h = W / 6.0, H / 4.0
+    lin = to_linear_srgb(warped)
+    means = []
+    for r in range(4):
+        for c in range(6):
+            x0 = int(c * cell_w + margin); x1 = int((c+1) * cell_w - margin)
+            y0 = int(r * cell_h + margin); y1 = int((r+1) * cell_h - margin)
+            roi = lin[y0:y1, x0:x1]
+            means.append(np.median(roi.reshape(-1, 3), axis=0))
+    return np.stack(means, 0)
+
+
+# ── calibration math ──────────────────────────────────────────────
+
+def white_balance_neutrals(src24, ref24):
+    idx = np.arange(18, 24)
+    src_g = src24[idx].mean(0)
+    ref_g = ref24[idx].mean(0)
+    gains = ref_g / np.maximum(src_g, 1e-6)
+    return gains
+
+
+def solve_ccm_no_bias(src24, ref24, use_indices):
+    A = src24[use_indices]
+    B = ref24[use_indices]
+    X, *_ = np.linalg.lstsq(A, B, rcond=None)
+    M = X.T
+    return M
+
+
+def fit_monotone_luma_curve_midgrays(src24_lin, ref24_lin):
+    mid_idx = np.array([19, 20, 21, 22])
+    w = np.array([0.2126, 0.7152, 0.0722], np.float32)
+
+    Ls = (src24_lin[mid_idx] @ w).astype(np.float32)
+    Lt = (ref24_lin[mid_idx] @ w).astype(np.float32)
+
+    eps_lo, eps_hi = 0.01, 0.98
+    Ls = np.concatenate([[eps_lo], np.sort(Ls), [eps_hi]])
+    Lt = np.concatenate([[eps_lo], np.sort(Lt), [eps_hi]])
+
+    def map_luma(L):
+        L = np.clip(L, 0, 1)
+        j = np.searchsorted(Ls, L, side='right') - 1
+        j = np.clip(j, 0, len(Ls)-2)
+        t = (L - Ls[j]) / np.maximum(Ls[j+1] - Ls[j], 1e-6)
+        return (1 - t) * Lt[j] + t * Lt[j+1]
+    return map_luma
+
+
+def soft_highlight_rolloff(L, knee=0.90, strength=0.6):
+    below = L < knee
+    out = np.empty_like(L, dtype=np.float32)
+    out[below] = L[below]
+    x = (L[~below] - knee) / max(1e-6, (1.0 - knee))
+    out[~below] = knee + (1.0 - knee) * (1.0 - np.exp(-strength * x))
+    return out
+
+
+# ── full correction pipeline ─────────────────────────────────────
+
+def apply_pipeline(target_bgr_u8, ref24, tgt24):
+    gains = white_balance_neutrals(tgt24, ref24)
+
+    lin = to_linear_srgb(target_bgr_u8)
+    lin_wb = lin * gains.reshape(1, 1, 3)
+
+    tgt24_wb = tgt24 * gains
+
+    chroma_idx = np.arange(0, 18)
+    M = solve_ccm_no_bias(tgt24_wb, ref24, chroma_idx)
+
+    H, W = lin_wb.shape[:2]
+    corrected = lin_wb.reshape(-1, 3) @ M.T
+    corrected = corrected.reshape(H, W, 3)
+
+    map_luma = fit_monotone_luma_curve_midgrays(tgt24_wb, ref24)
+    w = np.array([0.2126, 0.7152, 0.0722], np.float32)
+    L = np.clip(np.tensordot(corrected, w, axes=([2], [0])), 0, 1)
+    Lt = map_luma(L)
+    Lt = soft_highlight_rolloff(Lt, knee=0.90, strength=0.6)
+    eps = 1e-6
+    scale = (Lt + eps) / (L + eps)
+    corrected = corrected * scale[..., None]
+
+    corrected = np.clip(corrected, 0, 1)
+    return to_srgb_u8(corrected)

test_color_eval.py

@@ -0,0 +1,122 @@
+import os
+import csv
+import cv2 as cv
+import numpy as np
+from skimage import color
+import test_color_calibration
+
+# ── 1. Patch-Level Pipeline Simulation ────────────────────────────
+def apply_pipeline_to_patches(tgt24_lin, ref24_lin):
+    """
+    Applies the exact same math as `apply_pipeline`, but strictly to the 
+    24 patch array. This is much faster for calculating metrics.
+    """
+    # 1. White Balance
+    gains = test_color_calibration.white_balance_neutrals(tgt24_lin, ref24_lin)
+    tgt24_wb = tgt24_lin * gains
+    
+    # 2. CCM
+    chroma_idx = np.arange(0, 18)
+    M = test_color_calibration.solve_ccm_no_bias(tgt24_wb, ref24_lin, chroma_idx)
+    tgt24_ccm = tgt24_wb @ M.T
+    
+    # 3. Luma Curve
+    map_luma = test_color_calibration.fit_monotone_luma_curve_midgrays(tgt24_ccm, ref24_lin)
+    w = np.array([0.2126, 0.7152, 0.0722], np.float32)
+    
+    L = np.clip(np.tensordot(tgt24_ccm, w, axes=([1], [0])), 0, 1)
+    Lt = map_luma(L)
+    Lt = test_color_calibration.soft_highlight_rolloff(Lt, knee=0.90, strength=0.6)
+    
+    eps = 1e-6
+    scale = (Lt + eps) / (L + eps)
+    tgt24_corrected = tgt24_ccm * scale[:, None]
+    
+    return np.clip(tgt24_corrected, 0, 1)
+
+# ── 2. Delta E Calculation ────────────────────────────────────────
+def compute_deltaE_00(lin_src, lin_ref):
+    """
+    Computes the CIEDE2000 color difference between two sets of linear RGB patches.
+    """
+    def lin_to_srgb(lin):
+        a = 0.055
+        srgb = np.where(lin <= 0.0031308, 12.92 * lin, (1 + a) * np.power(np.maximum(lin, 0), 1/2.4) - a)
+        return np.clip(srgb, 0, 1)
+        
+    # Convert linear RGB to standard sRGB float (0-1)
+    srgb_src = lin_to_srgb(lin_src).reshape(1, -1, 3)
+    srgb_ref = lin_to_srgb(lin_ref).reshape(1, -1, 3)
+    
+    # Convert sRGB to CIELAB color space
+    lab_src = color.rgb2lab(srgb_src)
+    lab_ref = color.rgb2lab(srgb_ref)
+    
+    # Calculate Delta E 2000
+    de = color.deltaE_ciede2000(lab_src, lab_ref)
+    return de.flatten() # Returns array of 24 Delta E values
+
+# ── 3. Dataset Batch Processor ────────────────────────────────────
+def quantify_dataset(reference_image_path, target_images_dir, output_csv="color_metrics.csv"):
+    # 1. Extract the Golden Reference patches
+    ref_img = cv.imread(reference_image_path)
+    if ref_img is None: raise FileNotFoundError(f"Could not load reference {reference_image_path}")
+    
+    ref_corners = test_color_calibration.detect_checker_corners(ref_img)
+    ref_warped = test_color_calibration.warp_checker(ref_img, ref_corners)
+    ref24 = test_color_calibration.sample_24_patches(ref_warped)
+
+    results =[]
+
+    # 2. Iterate through the dataset
+    for filename in sorted(os.listdir(target_images_dir)):
+        if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.tif')):
+            continue
+            
+        filepath = os.path.join(target_images_dir, filename)
+        tgt_img = cv.imread(filepath)
+        
+        try:
+            # Extract target patches
+            tgt_corners = test_color_calibration.detect_checker_corners(tgt_img)
+            tgt_warped = test_color_calibration.warp_checker(tgt_img, tgt_corners)
+            tgt24 = test_color_calibration.sample_24_patches(tgt_warped)
+            
+            # Metric 1: Initial Error (Before Calibration)
+            de_initial = compute_deltaE_00(tgt24, ref24)
+            
+            # Apply mathematical correction to patches
+            tgt24_corrected = apply_pipeline_to_patches(tgt24, ref24)
+            
+            # Metric 2: Final Error (After Calibration)
+            de_final = compute_deltaE_00(tgt24_corrected, ref24)
+            
+            # Save metrics
+            results.append({
+                'Image': filename,
+                'Initial_Mean_dE': np.mean(de_initial),
+                'Initial_Max_dE': np.max(de_initial),
+                'Final_Mean_dE': np.mean(de_final),
+                'Final_Max_dE': np.max(de_final),
+                'Final_Gray_dE': np.mean(de_final[18:24]), # Accuracy of grays specifically
+                'Final_Color_dE': np.mean(de_final[0:18])  # Accuracy of chroma specifically
+            })
+            print(f"Processed {filename}: dE improved from {np.mean(de_initial):.2f} -> {np.mean(de_final):.2f}")
+            
+        except Exception as e:
+            print(f"Skipping {filename} - Error: {e}")
+
+    # 3. Save to CSV
+    if results:
+        keys = results[0].keys()
+        with open(output_csv, 'w', newline='') as output_file:
+            dict_writer = csv.DictWriter(output_file, keys)
+            dict_writer.writeheader()
+            dict_writer.writerows(results)
+        print(f"\nSaved metrics for {len(results)} images to {output_csv}")
+    
+    return results
+
+if __name__ == "__main__":
+    # Example Usage:
+    quantify_dataset("reference.png", "full_data/")