Update app.py

app.py

@@ -78,7 +78,8 @@ def align_to_reference(img: np.ndarray, ref: np.ndarray, max_size: int = 1200) -
     img_gray = cv2.cvtColor(img_s, cv2.COLOR_BGR2GRAY)
     ref_gray = cv2.cvtColor(ref_s, cv2.COLOR_BGR2GRAY)
 
-    orb = cv2.ORB_create(5000)
+    # Increased keypoints for better matching on tyre textures
+    orb = cv2.ORB_create(10000)  # was 5000
     k1, d1 = orb.detectAndCompute(img_gray, None)
     k2, d2 = orb.detectAndCompute(ref_gray, None)
 
@@ -89,8 +90,8 @@ def align_to_reference(img: np.ndarray, ref: np.ndarray, max_size: int = 1200) -
     matches = bf.match(d1, d2)
     matches = sorted(matches, key=lambda x: x.distance)
 
-    if len(matches) < 20:
-        # not enough matches; naive resize to ref
+    # Require more matches before trusting homography
+    if len(matches) < 40:  # was 20
         warped = cv2.resize(img, (ref.shape[1], ref.shape[0]))
         return AlignResult(False, warped, None, len(matches))
 
@@ -119,10 +120,23 @@ def build_reference(images: List[np.ndarray]) -> np.ndarray:
     return ref
 
 
+# -------- New: illumination normalization helpers --------
+def match_mean_std(src_gray: np.ndarray, tgt_gray: np.ndarray) -> np.ndarray:
+    """Match mean/std of src to tgt to reduce lighting-induced SSIM differences."""
+    s = src_gray.astype(np.float32); t = tgt_gray.astype(np.float32)
+    sm, ss = float(s.mean()), float(s.std() + 1e-6)
+    tm, ts = float(t.mean()), float(t.std() + 1e-6)
+    out = ((s - sm) / ss) * ts + tm
+    return np.clip(out, 0, 255).astype(np.uint8)
+
+
 def ssim_anomaly_map(img: np.ndarray, ref: np.ndarray, win_size: int = 11, gaussian_weights: bool = True) -> np.ndarray:
-    """Return anomaly map = 1 - SSIM per-pixel (grayscale)."""
+    """Return anomaly map = 1 - SSIM per-pixel (grayscale) with illumination normalization."""
     img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     ref_gray = cv2.cvtColor(ref, cv2.COLOR_BGR2GRAY)
+    # Normalize illumination
+    img_gray = match_mean_std(img_gray, ref_gray)
+
     img_blur = cv2.GaussianBlur(img_gray, (5,5), 0)
     ref_blur = cv2.GaussianBlur(ref_gray, (5,5), 0)
     score, ssim_map = ssim(ref_blur, img_blur, win_size=win_size, gaussian_weights=gaussian_weights, full=True)
@@ -143,6 +157,17 @@ def postprocess_mask(anom_map: np.ndarray, thr: float, min_area: int) -> np.ndar
     return mask_bool.astype(np.uint8)
 
 
+# -------- New: tyre ROI mask to ignore background/edges --------
+def center_circular_mask(shape, keep_ratio=0.9) -> np.ndarray:
+    """Return a central circular mask to focus on the tyre and ignore background edge shifts."""
+    h, w = shape[:2]
+    cy, cx = h//2, w//2
+    r = int(min(h, w) * keep_ratio / 2)
+    Y, X = np.ogrid[:h, :w]
+    mask = ((Y - cy)**2 + (X - cx)**2) <= r*r
+    return mask.astype(np.uint8)
+
+
 def overlay_and_boxes(img: np.ndarray, mask: np.ndarray) -> Tuple[np.ndarray, list]:
     overlay = img.copy()
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -157,8 +182,6 @@ def overlay_and_boxes(img: np.ndarray, mask: np.ndarray) -> Tuple[np.ndarray, li
     # red translucent overlay
     red = np.zeros_like(img)
     red[:] = (0, 0, 255)
-    alpha = (mask * 120)[:, :, None]  # 0..255 like weight
-    alpha = alpha.astype(np.uint8)
     heat = cv2.addWeighted(img, 1.0, red, 0.35, 0)
     out = np.where(mask[:, :, None].astype(bool), heat, img)
     # draw boxes on top
@@ -173,6 +196,16 @@ def mask_stats(mask: np.ndarray) -> dict:
     return {"defect_area_px": area, "image_area_px": total, "defect_pct": 100.0 * area / total}
 
 
+# -------- New: global SSIM guard --------
+def global_ssim(img: np.ndarray, ref: np.ndarray) -> float:
+    g1 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    g2 = cv2.cvtColor(ref, cv2.COLOR_BGR2GRAY)
+    g1 = cv2.GaussianBlur(g1, (5,5), 0)
+    g2 = cv2.GaussianBlur(g2, (5,5), 0)
+    score, _ = ssim(g2, g1, full=True)
+    return float(score)
+
+
 # ---------------------------
 # Streamlit UI
 # ---------------------------
@@ -188,17 +221,22 @@ with st.sidebar:
 
     st.header("Detection Settings")
     align = st.checkbox("Auto-align to reference (ORB)", value=True)
-    win = st.select_slider("SSIM window", options=[7,9,11,13], value=11)
-    thr = st.slider("Anomaly threshold (1-SSIM)", 0.05, 0.80, 0.35, step=0.01)
-    min_area = st.slider("Min defect area (px)", 10, 5000, 300, step=10)
-    decision_pct = st.slider("Fail if defect area > (%)", 0.1, 10.0, 1.0, step=0.1)
+    apply_roi = st.checkbox("Apply tyre ROI mask (ignore background/edges)", value=True)
+    win = st.select_slider("SSIM window", options=[7,9,11,13], value=13)  # default relaxed
+    thr = st.slider("Anomaly threshold (1-SSIM)", 0.05, 0.80, 0.45, step=0.01)  # default higher
+    min_area = st.slider("Min defect area (px)", 10, 10000, 1000, step=10)  # default higher
+    decision_pct = st.slider("Fail if defect area > (%)", 0.1, 10.0, 2.5, step=0.1)  # default higher
 
 if 'reference' not in st.session_state:
     st.session_state.reference = None
 
 if build_btn and ref_files:
     imgs = [resize_max(read_image(f)) for f in ref_files]
-    ref = build_reference(imgs)
+    # If only one good image is provided, use it directly to avoid median-shift
+    if len(imgs) == 1:
+        ref = imgs[0]
+    else:
+        ref = build_reference(imgs)
     st.session_state.reference = ref
     st.success(f"Reference built from {len(imgs)} image(s). Size: {ref.shape[1]}×{ref.shape[0]}")
 
@@ -208,7 +246,7 @@ if ref is None:
     st.info("No reference yet. Upload 1–10 good images in the sidebar and click *Build / Update Reference*.")
 else:
     st.subheader("1) Golden Reference Preview")
-    st.image(cv_to_pil(ref), caption="Golden Reference (median composite)", use_column_width=True)
+    st.image(cv_to_pil(ref), caption="Golden Reference (median composite or single good image)", use_column_width=True)
 
 st.markdown("---")
 
@@ -241,13 +279,24 @@ if ref is not None and test_files:
         if img is None:
             st.warning(f"Could not read {getattr(f,'name','file')} — skipping.")
             continue
+
         A = align_to_reference(img, ref) if align else AlignResult(False, cv2.resize(img, (ref.shape[1], ref.shape[0])), None, 0)
         warped = A.warped
+
         anom = ssim_anomaly_map(warped, ref, win_size=win)
+        # Optionally restrict to tyre ROI
+        if apply_roi:
+            roi = center_circular_mask(anom.shape, keep_ratio=0.9)
+            anom = anom * roi
+
         mask = postprocess_mask(anom, thr=thr, min_area=min_area)
         overlay, boxes = overlay_and_boxes(warped, mask)
         stats = mask_stats(mask)
-        label = "Defective" if stats['defect_pct'] > decision_pct else "Good"
+
+        # Global SSIM guard: if overall match is very high, prefer "Good"
+        gssim = global_ssim(warped, ref)
+        decision_is_defect = (stats['defect_pct'] > decision_pct) and (gssim < 0.985)
+        label = "Defective" if decision_is_defect else "Good"
 
         # UI
         st.write(f"*File:* {getattr(f,'name','image')}  |  Matches: {A.matches}  |  Alignment: {'OK' if A.ok else 'Fallback'}")
@@ -258,6 +307,7 @@ if ref is not None and test_files:
             st.metric("Decision", label)
             st.metric("Defect area %", f"{stats['defect_pct']:.2f}%")
             st.metric("Blobs detected", len(boxes))
+            st.metric("Global SSIM", f"{gssim:.4f}")
 
         # Prepare downloadable annotated image
         _ok, png_bytes = cv2.imencode('.png', overlay)
@@ -271,6 +321,7 @@ if ref is not None and test_files:
             "alignment_matches": A.matches,
             "aligned": A.ok,
             **stats,
+            "global_ssim": gssim,
             "decision": label,
             "boxes": len(boxes)
         })
@@ -284,4 +335,4 @@ else:
     st.info("Upload test images to run detection, or click *Generate a Demo from Reference* if a reference exists.")
 
 st.markdown("---")
-st.caption("©️ 2025 Procelevate Consulting — Demo app. For internal demonstration only; camera pose and lighting consistency are required.")
+st.caption("©️ 2025 Procelevate Consulting — Demo app. For internal demonstration only; camera pose and lighting consistency are required.")