Upload app.py

app.py

@@ -59,93 +59,85 @@ def detect_reference(
     mode: str,
     ref_size_mm: Optional[float],
 ) -> Tuple[float, Optional[Tuple[int, int]], Optional[str], Optional[Tuple[int, int, int, int]]]:
-    """Detect reference object (coin / square) in top-left ROI.
+    """Detect reference object (circle or square) using connected components.
 
-    Parameters
-    ----------
-    img_bgr : np.ndarray
-        Input BGR image.
-    mode : {"auto", "coin", "square"}
-        Detection strategy.
-    ref_size_mm : float or None
-        Real-world size (diameter for coin / side length for square).
-
-    Returns
-    -------
-    px_per_mm : float
-        Pixels per millimeter. Always > 0.
-    center : (int, int) or None
-        Reference object center in image coordinates.
-    ref_type : str or None
-        "coin" or "square" if detected, otherwise None.
+    Assumptions:
+    - White or near-white uniform background
+    - A single reference object is placed in the top-left region
+    - Reference is approximately square in its bounding box (square card or coin)
+    - ref_size_mm is the real diameter (coin) or side length (square)
     """
     h, w = img_bgr.shape[:2]
 
-    # Use a ROI at the top-left to limit search cost
-    roi_w = int(w * 0.25)
-    roi_h = int(h * 0.25)
-    roi = img_bgr[0:roi_h, 0:roi_w]
+    # 1. Estimate background color in LAB space and build "non-background" mask
+    #    This works for any solid-color background as long as the reference object
+    #    has a noticeable color difference from the background.
+    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
+    # Use the median color of the whole image as background estimate (robust to small objects)
+    bg_color = np.median(lab.reshape(-1, 3), axis=0)
+
+    # Compute per-pixel Euclidean distance in LAB space
+    diff = lab - bg_color  # shape (H, W, 3)
+    dist = np.sqrt(np.sum(diff * diff, axis=2)).astype(np.float32)
+
+    # Threshold on color distance: pixels far from background color are foreground
+    # You can tune 8.0 -> 6.0 or 10.0 depending on image contrast.
+    _, mask = cv2.threshold(dist, 8.0, 255, cv2.THRESH_BINARY)
+    mask = mask.astype(np.uint8)
+
+    # 2. Small morphological opening to remove noise
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
+
+    # 3. Connected components
+    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask)
+
+    # stats[i] = [x, y, w, h, area]
+    candidates = []
+    for i in range(1, num_labels):  # skip label 0 (background)
+        x, y, ww, hh, area = stats[i]
+        if area < 400:
+            # too small, likely noise
+            continue
+
+        # Only consider objects in the upper-left region
+        if x > w * 0.6 or y > h * 0.6:
+            continue
+
+        # Require roughly square bounding box: circles and squares both satisfy this
+        ar = ww / float(hh + 1e-6)
+        if ar < 0.7 or ar > 1.3:
+            continue
 
-    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
-    gray = cv2.medianBlur(gray, 5)
+        candidates.append((i, x, y, ww, hh, area))
 
-    px_per_mm: Optional[float] = None
+    px_per_mm: float
     center: Optional[Tuple[int, int]] = None
     ref_type: Optional[str] = None
     bbox: Optional[Tuple[int, int, int, int]] = None
 
-    # ----------------- coin detection -----------------
-    if mode in ("auto", "coin"):
-        circles = cv2.HoughCircles(
-            gray,
-            cv2.HOUGH_GRADIENT,
-            dp=1.2,
-            minDist=20,
-            param1=120,
-            param2=35,
-            minRadius=8,
-            maxRadius=min(roi_w, roi_h) // 2,
-        )
-        if circles is not None and len(circles) > 0:
-            c = circles[0][0]
-            r = float(c[2])
-            d_px = 2.0 * r
-            d_mm = ref_size_mm if ref_size_mm and ref_size_mm > 0 else 25.0
-            px_per_mm = max(d_px / d_mm, 1e-6)
-            center = (int(c[0]), int(c[1]))
-            ref_type = "coin"
-            bbox = (int(c[0] - r), int(c[1] - r), int(2 * r), int(2 * r))
-
-    # ----------------- square detection -----------------
-    if px_per_mm is None and mode in ("auto", "square"):
-        edges = cv2.Canny(gray, 80, 160)
-        cnts, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        best = None
-        best_score = 0.0
-        for cnt in cnts:
-            x, y, ww, hh = cv2.boundingRect(cnt)
-            area = ww * hh
-            if area < 225:
-                continue
-            score = min(ww, hh) / max(ww, hh)
-            if score > best_score:
-                best_score = score
-                best = (x, y, ww, hh)
-        if best is not None and best_score > 0.6:
-            x, y, ww, hh = best
-            s_px = float(max(ww, hh))
-            s_mm = ref_size_mm if ref_size_mm and ref_size_mm > 0 else 20.0
-            px_per_mm = max(s_px / s_mm, 1e-6)
-            center = (x + ww // 2, y + hh // 2)
-            ref_type = "square"
-            bbox = (x, y, ww, hh)
+    if candidates:
+        # 4. Pick the one closest to the top-left corner (smallest x + y)
+        label_id, x, y, ww, hh, area = min(candidates, key=lambda t: t[1] + t[2])
+        bbox = (int(x), int(y), int(ww), int(hh))
+        center = (int(x + ww // 2), int(y + hh // 2))
+
+        # Real-world size: diameter (coin) or side length (square)
+        ref_mm = ref_size_mm if ref_size_mm and ref_size_mm > 0 else 20.0
+
+        # For both circles and squares, the max side of the bounding box
+        # can be treated as "diameter/side" in pixels.
+        side_or_diam_px = float(max(ww, hh))
+        px_per_mm = max(side_or_diam_px / ref_mm, 1e-6)
 
-    # ----------------- fallback -----------------
-    if px_per_mm is None:
-        # Fallback: approximate scale if no reference detected.
-        # For typical scanner/phone images this is a safe range.
-        # Use a slightly conservative default so values不会太夸张.
+        # Roughly classify reference type; optional, not used in scaling
+        ref_type = "square"
+    else:
+        # If no reference found, use a safe default scale to avoid division by zero.
         px_per_mm = 4.0
+        center = None
+        ref_type = None
+        bbox = None
 
     return px_per_mm, center, ref_type, bbox
 
@@ -394,10 +386,10 @@ def analyze(
     color_tol: int,
     hsv_low_h: int,
     hsv_high_h: int,
-) -> Tuple[Optional[np.ndarray], pd.DataFrame, Optional[str], List[Dict[str, Any]]]:
+) -> Tuple[Optional[np.ndarray], pd.DataFrame, Optional[str], List[Dict[str, Any]], Dict[str, Any]]:
     try:
         if image is None:
-            return None, pd.DataFrame(), None, []
+            return None, pd.DataFrame(), None, [], {}
         img_rgb = np.array(image)
         img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
         img_bgr, scale = downscale_bgr(img_bgr)
@@ -416,13 +408,149 @@ def analyze(
         tmp.write(csv.encode("utf-8"))
         tmp.close()
         js = df.to_dict(orient="records")
-        return overlay, df, tmp.name, js
+
+        # Store state for interactive correction
+        state_dict: Dict[str, Any] = {
+            "img_bgr": img_bgr,
+            "sample_type": sample_type,
+            "px_per_mm": px_per_mm,
+            "ref_center": ref_center,
+            "ref_type": ref_type,
+            "ref_bbox": ref_bbox,
+            "components": comps,
+            "expected_count": expected_count,
+            "ref_size_mm": ref_size_mm,
+        }
+        # By default, all components are active samples
+        state_dict["active_indices"] = list(range(len(comps)))
+
+        return overlay, df, tmp.name, js, state_dict
     except Exception as e:
-        return None, pd.DataFrame(), None, [{"error": str(e)}]
+        return None, pd.DataFrame(), None, [{"error": str(e)}], {}
+
+
+# --- Interactive correction helper ---
+def apply_corrections(
+    click_event,
+    state_dict: Dict[str, Any],
+    correction_mode: str,
+) -> Tuple[Dict[str, Any], Optional[np.ndarray], pd.DataFrame, Optional[str], List[Dict[str, Any]]]:
+    """
+    Apply interactive corrections based on a click on the annotated image.
+
+    correction_mode:
+      - "none": do nothing
+      - "set-ref": treat the clicked object as the new reference
+      - "toggle-sample": toggle the clicked object between active/inactive sample
+    """
+    # If no valid state or no correction requested, do nothing
+    if not state_dict or "img_bgr" not in state_dict or correction_mode == "none" or click_event is None:
+        return state_dict, None, pd.DataFrame(), None, []
+
+    try:
+        # Gradio SelectData usually provides (x, y) in .index
+        if hasattr(click_event, "index"):
+            x, y = click_event.index
+        else:
+            # Fallback: assume click_event is a tuple
+            x, y = click_event
+
+        img_bgr = state_dict["img_bgr"]
+        components: List[Dict[str, Any]] = state_dict.get("components", [])
+        if not components:
+            return state_dict, None, pd.DataFrame(), None, []
+
+        # Find nearest component center to the click
+        min_dist = 1e9
+        nearest_idx = -1
+        for i, comp in enumerate(components):
+            cx, cy = comp["center"]
+            d = (cx - x) ** 2 + (cy - y) ** 2
+            if d < min_dist:
+                min_dist = d
+                nearest_idx = i
+
+        if nearest_idx < 0:
+            return state_dict, None, pd.DataFrame(), None, []
+
+        px_per_mm = state_dict.get("px_per_mm", 4.0)
+        ref_center = state_dict.get("ref_center")
+        ref_type = state_dict.get("ref_type", "square")
+        ref_bbox = state_dict.get("ref_bbox")
+        ref_size_mm = state_dict.get("ref_size_mm", 20.0)
+        sample_type = state_dict.get("sample_type", "leaves")
+
+        active_indices = state_dict.get("active_indices", list(range(len(components))))
+
+        if correction_mode == "set-ref":
+            # Use this component as the new reference object
+            comp = components[nearest_idx]
+            box = comp["box"]
+            xs = box[:, 0]
+            ys = box[:, 1]
+            x0, y0 = int(xs.min()), int(ys.min())
+            w0, h0 = int(xs.max() - xs.min()), int(ys.max() - ys.min())
+            ref_bbox = (x0, y0, w0, h0)
+            ref_center = (int(comp["center"][0]), int(comp["center"][1]))
+
+            # Update px_per_mm using the largest side as diameter/side length
+            side_px = float(max(w0, h0))
+            px_per_mm = max(side_px / (ref_size_mm if ref_size_mm > 0 else 20.0), 1e-6)
+            ref_type = "square"
+
+            # Remove this component from active samples (reference is not a sample)
+            new_components = []
+            for i, c in enumerate(components):
+                if i != nearest_idx:
+                    new_components.append(c)
+            components = new_components
+            # Rebuild active_indices to cover all remaining components
+            active_indices = list(range(len(components)))
+
+            state_dict["components"] = components
+            state_dict["ref_bbox"] = ref_bbox
+            state_dict["ref_center"] = ref_center
+            state_dict["px_per_mm"] = px_per_mm
+            state_dict["ref_type"] = ref_type
+            state_dict["active_indices"] = active_indices
+
+        elif correction_mode == "toggle-sample":
+            # Toggle this component in/out of the active sample set
+            if nearest_idx in active_indices:
+                active_indices = [idx for idx in active_indices if idx != nearest_idx]
+            else:
+                active_indices.append(nearest_idx)
+                active_indices = sorted(set(active_indices))
+            state_dict["active_indices"] = active_indices
+
+        # Rebuild the list of active components
+        active_components = [components[i] for i in active_indices]
+
+        # Recompute metrics and overlay using the updated state
+        df = compute_metrics(img_bgr, active_components, px_per_mm)
+        overlay = render_overlay(
+            img_bgr.copy(),
+            px_per_mm,
+            (state_dict.get("ref_center"), state_dict.get("ref_type")),
+            active_components,
+            df,
+            state_dict.get("ref_bbox"),
+        )
+        csv = df.to_csv(index=False)
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
+        tmp.write(csv.encode("utf-8"))
+        tmp.close()
+        js = df.to_dict(orient="records")
+
+        return state_dict, overlay, df, tmp.name, js
+    except Exception:
+        # In case of any error, do not break the app; just keep current state
+        return state_dict, None, pd.DataFrame(), None, []
 
 
 with gr.Blocks(theme=gr.themes.Default()) as demo:
     gr.Markdown("# Biological Sample Quantifier (Leaves / Seeds)")
+    state = gr.State({})
     with gr.Row():
         with gr.Column(scale=1):
             image = gr.Image(type="numpy", label="Upload image")
@@ -435,17 +563,48 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
             color_tol = gr.Slider(5, 100, value=40, step=1, label="Color tolerance")
             hsv_low = gr.Slider(0, 179, value=35, step=1, label="HSV H lower (leaves)")
             hsv_high = gr.Slider(0, 179, value=85, step=1, label="HSV H upper (leaves)")
+            correction_mode = gr.Radio(
+                ["none", "set-ref", "toggle-sample"],
+                value="none",
+                label="Correction mode (click on image)"
+            )
             run = gr.Button("Analyze")
             reset = gr.Button("Reset")
         with gr.Column(scale=2):
-            overlay = gr.Image(label="Annotated")
+            overlay = gr.Image(label="Annotated", interactive=True)
             table = gr.Dataframe(label="Metrics", wrap=True)
             csv_out = gr.File(label="CSV export")
             json_out = gr.JSON(label="JSON preview")
     def _analyze(image, sample_type, expected, ref_mode, ref_size, min_area, max_area, color_tol, hsv_low, hsv_high):
-        return analyze(image, sample_type, expected, ref_mode, ref_size, min_area, max_area, color_tol, hsv_low, hsv_high)
-    run.click(_analyze, [image, sample_type, expected, ref_mode, ref_size, min_area, max_area, color_tol, hsv_low, hsv_high], [overlay, table, csv_out, json_out])
-    reset.click(lambda: (None, pd.DataFrame(), None, []), None, [overlay, table, csv_out, json_out])
+        overlay_img, df, csv_path, js, state_dict = analyze(
+            image, sample_type, expected, ref_mode, ref_size, min_area, max_area, color_tol, hsv_low, hsv_high
+        )
+        return overlay_img, df, csv_path, js, state_dict
+
+    run.click(
+        _analyze,
+        [image, sample_type, expected, ref_mode, ref_size, min_area, max_area, color_tol, hsv_low, hsv_high],
+        [overlay, table, csv_out, json_out, state],
+    )
+
+    def _reset():
+        return None, pd.DataFrame(), None, [], {}
+
+    reset.click(_reset, None, [overlay, table, csv_out, json_out, state])
+
+    def _on_select(evt, current_state, correction_mode):
+        # Apply corrections based on a click on the annotated image
+        new_state, overlay_img, df, csv_path, js = apply_corrections(evt, current_state or {}, correction_mode)
+        # If overlay_img is None, keep the existing outputs unchanged by returning gr.update()
+        if overlay_img is None:
+            return gr.update(), gr.update(), gr.update(), gr.update(), new_state
+        return overlay_img, df, csv_path, js, new_state
+
+    overlay.select(
+        _on_select,
+        [state, correction_mode],
+        [overlay, table, csv_out, json_out, state],
+    )
 
 if __name__ == "__main__":
     demo.launch()