Update app.py

app.py

@@ -9,42 +9,143 @@ from PIL import Image, ImageDraw
 import io
 from huggingface_hub import hf_hub_download
 import base64
-from concurrent.futures import ThreadPoolExecutor, as_completed
 import csv
 import joblib
 import os
-
-HF_REPO_ID = "myang4218/cellposemodel"
-HF_REPO_ID2 = "LiangLabUMB/viability_model"
+import sys
+import json
+
+from error_labeling_tab import build_error_labeling_tab
+
+# error_detection/'s modules (correct.py, inference.py, model.py, ...) use bare
+# intra-package imports (e.g. "from dataset import ...") and assume their own
+# directory is on sys.path, per how their test suite is run
+# (track_a/TrackA_SelfCorrection_ImplementationGuide.md, Step 7-9).
+_ERROR_DETECTION_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "track_a", "error_detection")
+if _ERROR_DETECTION_DIR not in sys.path:
+    sys.path.insert(0, _ERROR_DETECTION_DIR)
+
+HF_REPO_ID    = "myang4218/cellposemodel"
+HF_REPO_ID2   = "LiangLabUMB/viability_model"
+HF_REPO_CPSAM = "mouseland/cellpose-sam"
 MODEL_OPTIONS = {
     "Hemocytometer Model": "hemocytometermodel.npy",
-    "General Model": "generalmodel.npy"
+    "General Model":       "generalmodel.npy",
+    "Cellpose SAMv2":      "cpsam_v2",
+}
+MODEL_REPOS = {
+    "hemocytometermodel.npy": HF_REPO_ID,
+    "generalmodel.npy":       HF_REPO_ID,
+    "cpsam_v2":               HF_REPO_CPSAM,
 }
 
-loaded_models = {}
-
-VIABILITY_CLF    = None
-VIABILITY_SCALER = None
- 
-try:
-    _clf_path    = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_clf.pkl")
-    _scaler_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_scaler.pkl")
-    VIABILITY_CLF    = joblib.load(_clf_path)
-    VIABILITY_SCALER = joblib.load(_scaler_path)
-    print("✓ Viability classifier loaded.")
-except Exception as e:
-    print(f"Viability classifier not found or failed to load: {e}")
-
-# ---- mobile-safe size limits (aggressive for Safari) ----
+
+# Viability classifier is loaded LAZILY (only when viability is first requested).
+# Loading XGBoost / the pickled model at module level initialises a CUDA context
+# in the main process, which conflicts with ZeroGPU's per-call GPU allocation and
+# leaves segmentation requests stuck in the GPU queue. Keeping the main process
+# CUDA-clean until viability is explicitly called avoids this.
+VIABILITY_CLF     = None
+VIABILITY_SCALER  = None
+_VIABILITY_LOADED = False
+_VIABILITY_ERROR  = None
+
+
+def ensure_viability_loaded():
+    """
+    Load the XGBoost classifier + scaler on first use, pinned to CPU.
+
+    XGBoost initialises a CUDA context if it thinks a GPU is available. Once
+    that context exists in the main process, every subsequent ZeroGPU fork for
+    run_segmentation inherits a dirty CUDA state and hangs in the GPU queue.
+    We force the loaded booster onto CPU so it never touches CUDA — viability
+    inference on ~20 features per cell is trivially fast on CPU anyway.
+    """
+    global VIABILITY_CLF, VIABILITY_SCALER, _VIABILITY_LOADED, _VIABILITY_ERROR
+    if _VIABILITY_LOADED:
+        return
+    try:
+        import xgboost
+        _clf_path    = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_xgb_clf.pkl")
+        _scaler_path = hf_hub_download(repo_id=HF_REPO_ID2, filename="viability_xgb_scaler.pkl")
+        VIABILITY_CLF    = joblib.load(_clf_path)
+        VIABILITY_SCALER = joblib.load(_scaler_path)
+
+        # Force CPU inference so XGBoost never initialises a CUDA context.
+        try:
+            VIABILITY_CLF.set_params(device="cpu", predictor="cpu_predictor", tree_method="hist")
+        except Exception:
+            pass
+        try:
+            booster = VIABILITY_CLF.get_booster()
+            booster.set_param({"device": "cpu", "predictor": "cpu_predictor"})
+        except Exception:
+            pass
+
+        _VIABILITY_LOADED = True
+        print("✓ Viability classifier loaded (lazy, CPU-pinned).")
+    except Exception as e:
+        _VIABILITY_ERROR = str(e)
+        print(f"Viability classifier failed to load: {e}")
+
+
+# Self-correction (Track A4, TrackA_SelfCorrection_ImplementationGuide.md Step
+# 7-9) model is loaded LAZILY, same rationale as the viability classifier:
+# keep the main process CUDA-clean for ZeroGPU. HybridErrorNet is a small
+# 5x96x96 CNN + MLP — batched CPU inference over every instance in an image is
+# fast, so this never needs @spaces.GPU.
+SELF_CORRECTION_MODEL   = None
+SELF_CORRECTION_TEMP    = 1.0
+SELF_CORRECTION_RUN_FN  = None
+_SELF_CORRECTION_LOADED = False
+_SELF_CORRECTION_ERROR  = None
+
+
+def ensure_self_correction_loaded():
+    global SELF_CORRECTION_MODEL, SELF_CORRECTION_TEMP, SELF_CORRECTION_RUN_FN
+    global _SELF_CORRECTION_LOADED, _SELF_CORRECTION_ERROR
+    if _SELF_CORRECTION_LOADED:
+        return
+    try:
+        import torch
+        from model import HybridErrorNet
+        from inference import run_self_correction
+
+        ckpt_path = os.path.join(_ERROR_DETECTION_DIR, "checkpoints", "best.pt")
+        ckpt = torch.load(ckpt_path, map_location="cpu")
+        model = HybridErrorNet()
+        model.load_state_dict(ckpt["model_state"])
+        model.eval()
+        SELF_CORRECTION_MODEL = model
+        SELF_CORRECTION_RUN_FN = run_self_correction
+
+        # Temperature from calibrate.py's --val-dir run against real_val/ (see
+        # calibration.json's n_val_examples_by_class for how much real data it
+        # was fit on — merge is the scarcest class, hence the extra-cautious
+        # default merge threshold in inference.DEFAULT_THRESHOLDS).
+        calib_path = os.path.join(_ERROR_DETECTION_DIR, "checkpoints", "calibration.json")
+        if os.path.exists(calib_path):
+            with open(calib_path) as f:
+                SELF_CORRECTION_TEMP = json.load(f)["temperature"]
+
+        _SELF_CORRECTION_LOADED = True
+        # Plain ASCII: a checkmark here previously raised UnicodeEncodeError on
+        # a cp1252 (default Windows console) stdout, which — since it ran
+        # inside this try block — got mis-caught below as a load failure even
+        # though the model had already loaded successfully.
+        print(f"Self-correction model loaded (lazy, CPU-pinned, T={SELF_CORRECTION_TEMP:.2f}).")
+    except Exception as e:
+        _SELF_CORRECTION_ERROR = str(e)
+        print(f"Self-correction model failed to load: {e}")
+
+
+# mobile safe resize limits
 MAX_SIDE = 1024          
 MAX_PIXELS = 1024 * 1024
 
 
 def safe_resize(image_np):
-    """
-    Downscale image to fit within MAX_SIDE and MAX_PIXELS while
-    preserving aspect ratio. Works for RGB / RGBA / grayscale.
-    """
+    
     h, w = image_np.shape[:2]
     total = h * w
 
@@ -152,11 +253,7 @@ FEATURE_COLS_INFERENCE = [
 
 
 def classify_cells_by_model(image_np, masks):
-    """
-    Run the trained LogisticRegression classifier to predict live/dead per cell.
-    Returns (dead_count, alive_count, overlay_np, {cell_id: label}).
-    Requires VIABILITY_CLF and VIABILITY_SCALER to be loaded.
-    """
+    
     import numpy as np
     cell_ids = np.unique(masks)
     cell_ids = cell_ids[cell_ids > 0]
@@ -187,12 +284,23 @@ def classify_cells_by_model(image_np, masks):
     return dead, alive, overlay, label_map
 
 
+def build_colored_mask_overlay(image_np, masks, alpha=0.4):
+    """Deterministic per-instance-ID random color blended over the image —
+    the exact rendering run_segmentation uses for its 'Segmentation Result'
+    view. Factored out here so the self-correction 'before'/'after' panels
+    can produce a visually consistent overlay from any given masks array."""
+    overlay = image_np.copy().astype(np.float32)
+    if masks.max() > 0:
+        np.random.seed(42)  # same seed as run_segmentation -> stable colors per ID
+        colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
+        colors[0] = [0, 0, 0]
+        colored_mask = colors[masks]
+        overlay = (1 - alpha) * overlay + alpha * colored_mask
+    return np.clip(overlay, 0, 255).astype(np.uint8)
+
+
 def draw_viability_overlay(image_np, masks, label_map):
-    """
-    Draw coloured contours + cell-number labels onto image_np.
-    label_map: {cell_id: 0=live, 1=dead}
-    Returns a uint8 numpy array.
-    """
+
     overlay  = image_np.copy()
     cell_ids = np.unique(masks)
     cell_ids = cell_ids[cell_ids > 0]
@@ -220,114 +328,10 @@ def draw_viability_overlay(image_np, masks, label_map):
                           (0, 0, 0), -1)
             cv2.putText(overlay, label_str,
                         (cx - tw//2, cy + th//2),
-                        font, font_scale, color, thickness, cv2.LINE_AA)
+                        font, font_scale, color, thickness, cv2.LINE_AA) 
     return overlay
 
 
-def classify_cells_by_blueness(image_np, masks, threshold_bias):
-    """
-    Classify cells as dead (blue) or alive using an adaptive Otsu threshold
-    on per-cell blueness scores, with a user bias to fine-tune.
-
-    Args:
-        image_np:        RGB image array
-        masks:           Cellpose segmentation masks
-        threshold_bias:  Slider value -50..+50; shifts Otsu threshold up/down.
-                         Negative = more cells classified dead (looser).
-                         Positive = fewer cells classified dead (stricter).
-                         0 = pure Otsu (fully automatic).
-
-    Returns:
-        dead_count, alive_count, colored_overlay, otsu_threshold, final_threshold
-    """
-
-    if len(image_np.shape) == 2:
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
-    elif len(image_np.shape) == 3 and image_np.shape[2] == 4:
-        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGBA2RGB)
-
-    hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)
-
-    hue        = hsv[:, :, 0].astype(np.float32)
-    saturation = hsv[:, :, 1].astype(np.float32)
-
-    # Raw blueness: hue proximity to 115° × saturation
-    hue_distance = np.minimum(np.abs(hue - 115), 180 - np.abs(hue - 115))
-    hue_score    = np.maximum(0, 1 - hue_distance / 65)
-    blueness     = hue_score * (saturation / 255.0)
-
-    # --- Compute per-cell mean blueness scores ---
-    cell_ids = np.unique(masks)
-    cell_ids = cell_ids[cell_ids > 0]
-
-    if len(cell_ids) == 0:
-        blank = image_np.copy()
-        return 0, 0, blank, 0.0, 0.0
-
-    cell_scores = np.array([np.mean(blueness[masks == cid]) for cid in cell_ids])
-
-    # --- Otsu on the distribution of per-cell scores ---
-    # cv2.threshold expects uint8; scale 0-1 → 0-255
-    scores_u8 = (np.clip(cell_scores, 0, 1) * 255).astype(np.uint8)
-
-    if scores_u8.max() == scores_u8.min():
-        # All cells identical → Otsu is undefined; use midpoint
-        otsu_threshold = float(scores_u8[0]) / 255.0
-    else:
-        # Reshape to a single-column image so cv2.threshold works
-        thresh_val, _ = cv2.threshold(
-            scores_u8.reshape(-1, 1), 0, 255,
-            cv2.THRESH_BINARY + cv2.THRESH_OTSU
-        )
-        otsu_threshold = thresh_val / 255.0
-
-    # --- Apply user bias: slider -50..+50 maps to ±0.20 shift ---
-    bias = (threshold_bias / 50.0) * 0.20
-    final_threshold = float(np.clip(otsu_threshold + bias, 0.0, 1.0))
-
-    # --- Classify ---
-    dead_cells  = [cid for cid, s in zip(cell_ids, cell_scores) if s > final_threshold]
-    alive_cells = [cid for cid, s in zip(cell_ids, cell_scores) if s <= final_threshold]
-
-    # --- Outline-only overlay on raw image with enumerated labels ---
-    final_overlay = image_np.copy()
-
-    # Compute a consistent enumeration order (cell_ids is already sorted ascending)
-    cell_enum = {cid: idx + 1 for idx, cid in enumerate(cell_ids)}
-
-    dead_set  = set(dead_cells)
-    alive_set = set(alive_cells)
-
-    for cid in cell_ids:
-        cell_mask = (masks == cid).astype(np.uint8)
-        contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        color = (220, 50, 50) if cid in dead_set else (50, 220, 80)
-        cv2.drawContours(final_overlay, contours, -1, color, thickness=2)
-
-        # Draw enumeration label at centroid
-        ys, xs = np.where(cell_mask)
-        if len(ys) > 0:
-            cx, cy = int(xs.mean()), int(ys.mean())
-            label_str = str(cell_enum[cid])
-            font       = cv2.FONT_HERSHEY_SIMPLEX
-            font_scale = 0.35
-            thickness  = 1
-            (tw, th), _ = cv2.getTextSize(label_str, font, font_scale, thickness)
-            # Dark background rectangle for readability
-            cv2.rectangle(
-                final_overlay,
-                (cx - tw // 2 - 1, cy - th // 2 - 1),
-                (cx + tw // 2 + 1, cy + th // 2 + 1),
-                (0, 0, 0),
-                -1
-            )
-            cv2.putText(
-                final_overlay, label_str,
-                (cx - tw // 2, cy + th // 2),
-                font, font_scale, color, thickness, cv2.LINE_AA
-            )
-
-    return len(dead_cells), len(alive_cells), final_overlay, otsu_threshold, final_threshold
 
 
 def measure_confluency(masks, image_np):
@@ -451,12 +455,10 @@ def warp_polygon_to_square(image_np, points):
 
 
 def toggle_stereological_mode(use_stereology):
-    """Show/hide stereological controls based on checkbox"""
     return gr.update(visible=use_stereology)
 
 
 def update_exclusion_preview(image, left_width, top_width):
-    """Update the preview image with exclusion zone overlay"""
     if image is None:
         return None
     
@@ -465,9 +467,8 @@ def update_exclusion_preview(image, left_width, top_width):
     return Image.fromarray(overlay)
 
 
-# ---------------------------------------------------------------------------
 # Patch segmentation
-# ---------------------------------------------------------------------------
+
 PATCH_SIZE   = 512   # target patch side length
 PATCH_OVERLAP = 64   # overlap border on each edge (pixels)
 MIN_PATCH_DIM = 256  # don't bother patching if image fits comfortably
@@ -555,62 +556,42 @@ def _merge_patch_masks(patch_results, full_h, full_w, overlap=PATCH_OVERLAP):
     return renumbered
 
 
-def _segment_patch(args):
-    """Worker: run cellpose on a single patch. Called from a thread pool."""
-    patch_np, row_start, col_start, model_filename, hf_repo = args
-    # Each thread uses the shared loaded_models cache (GIL-safe for reads;
-    # model.eval() releases the GIL during GPU work so threads overlap.)
-    model_path = hf_hub_download(repo_id=hf_repo, filename=model_filename)
-    if model_filename in loaded_models:
-        model = loaded_models[model_filename]
-    else:
-        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-        loaded_models[model_filename] = model
-
-    mask, _, _ = model.eval(patch_np, diameter=None, channels=[0, 0])
-    return mask, row_start, col_start
-
-
 def run_segmentation_patched(image_np, model_filename):
     """
-    Split image into overlapping patches, run Cellpose on each in parallel,
+    Split image into overlapping patches, run Cellpose sequentially on each,
     then stitch back into a single full-resolution mask.
-    Falls back to whole-image segmentation if the image is small enough
-    that patching adds overhead without benefit.
+
+    NOTE: ThreadPoolExecutor was removed because ZeroGPU allocates the GPU
+    only to the main thread. Worker threads spawned by ThreadPoolExecutor run
+    outside the GPU context, causing a 3-minute timeout waiting for a GPU
+    grant that never arrives for the worker threads.
+    Sequential patching within the @spaces.GPU context is correct and fast.
     """
     h, w = image_np.shape[:2]
-    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
-    if model_filename in loaded_models:
-        model = loaded_models[model_filename]
-    else:
-        model = models.CellposeModel(gpu=True, pretrained_model=model_path)
-        loaded_models[model_filename] = model
+    repo       = MODEL_REPOS.get(model_filename, HF_REPO_ID)
+    # hf_hub_download caches the weights file on disk, so this only hits the
+    # network once — subsequent calls read from the local cache.
+    model_path = hf_hub_download(repo_id=repo, filename=model_filename)
+
+    # IMPORTANT: build the CellposeModel fresh on every call. ZeroGPU allocates
+    # a DIFFERENT physical GPU for each @spaces.GPU invocation, so a model cached
+    # from a previous call holds a CUDA context bound to a now-deallocated GPU.
+    # Reusing it makes the next request hang in the GPU queue. Constructing the
+    # model is cheap because the weights are already on local disk.
+    model = models.CellposeModel(gpu=True, pretrained_model=model_path)
 
     # Small images: no benefit from patching
     if max(h, w) <= MIN_PATCH_DIM * 2:
-        mask, _, _ = model.eval(image_np, diameter=None, channels=[0, 0])
-        return mask, 1   # 1 patch
+        mask, _, _ = model.eval(image_np, diameter=None)
+        return mask, 1
 
     patches = _split_patches(image_np)
     n_patches = len(patches)
 
-    # Build argument list for the thread pool
-    args_list = [
-        (patch, r, c, model_filename, HF_REPO_ID)
-        for patch, r, c in patches
-    ]
-
-    patch_results = []  # (mask, row_start, col_start) in submission order
-
-    # ThreadPoolExecutor: GPU kernels release the GIL so threads overlap on GPU
-    with ThreadPoolExecutor(max_workers=min(n_patches, 4)) as pool:
-        futures = {pool.submit(_segment_patch, a): a for a in args_list}
-        for future in as_completed(futures):
-            mask_patch, row_start, col_start = future.result()
-            patch_results.append((mask_patch, row_start, col_start))
-
-    # Re-sort by (row, col) so stitching is deterministic
-    patch_results.sort(key=lambda x: (x[1], x[2]))
+    patch_results = []
+    for patch, row_start, col_start in patches:
+        mask_patch, _, _ = model.eval(patch, diameter=None)
+        patch_results.append((mask_patch, row_start, col_start))
 
     full_mask = _merge_patch_masks(patch_results, h, w)
     return full_mask, n_patches
@@ -618,6 +599,7 @@ def run_segmentation_patched(image_np, model_filename):
 
 @spaces.GPU
 def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
+                     use_min_filter, use_max_filter,
                      use_stereology, left_exclusion, top_exclusion,
                      crop_points=None):
     image_np = np.array(image)
@@ -661,21 +643,18 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
         print("p90:", np.percentile(sizes, 90) if len(sizes) > 0 else 0)
         print("max:", sizes.max() if len(sizes) > 0 else 0)
         
-        # Compute recommendation from RAW masks 
+        # Compute recommendation from RAW masks (always shown, never auto-applied)
         recommend_min = rec_min_size(masks_raw)
 
-        # If user sets slider to 0, use the recommendation
-        min_used = recommend_min if (min_cell_size == 0) else int(min_cell_size)
-
-        # Apply filters
+        # Apply filters only if their checkboxes are enabled
         masks = masks_raw.copy()
         removed_small = 0
         removed_large = 0
 
-        if min_used > 0:
-            masks, removed_small = filter_mask_by_size(masks, min_used)
+        if use_min_filter and int(min_cell_size) > 0:
+            masks, removed_small = filter_mask_by_size(masks, int(min_cell_size))
 
-        if max_cell_size > 0:
+        if use_max_filter and max_cell_size > 0:
             masks, removed_large = filter_mask_by_maxsize(masks, int(max_cell_size))
 
         # Apply stereological exclusion if enabled
@@ -687,7 +666,7 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
         
         filter_msg = ""
         if removed_small:
-            filter_msg += f"Removed {removed_small} small objects (< {min_used} pixels).\n"
+            filter_msg += f"Removed {removed_small} small objects (< {int(min_cell_size)} pixels).\n"
         if removed_large:
             filter_msg += f"Removed {removed_large} large objects (> {int(max_cell_size)} pixels).\n"
         if use_stereology and excluded_count > 0:
@@ -697,16 +676,8 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
         confluency = measure_confluency(masks, processed_image_np)
 
         # Create a basic segmentation overlay (without viability)
-        segmentation_overlay = processed_image_np.copy().astype(np.float32)
-        if masks.max() > 0:
-            np.random.seed(42)  # For consistent random colors
-            colors = np.random.randint(0, 255, size=(masks.max() + 1, 3))
-            colors[0] = [0, 0, 0]
-            colored_mask = colors[masks]
-            alpha = 0.4
-            segmentation_overlay = (1 - alpha) * segmentation_overlay + alpha * colored_mask
-        segmentation_overlay = np.clip(segmentation_overlay, 0, 255).astype(np.uint8)
-        
+        segmentation_overlay = build_colored_mask_overlay(processed_image_np, masks)
+
         # Add exclusion zone overlay if stereology is enabled
         if use_stereology:
             segmentation_overlay = draw_exclusion_overlay(segmentation_overlay, left_exclusion, top_exclusion)
@@ -729,7 +700,7 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
             pack_array(masks),
             pack_array(processed_image_np),
             confluency,
-            gr.update(value=recommend_min), 
+            f"Recommended minimum: **{recommend_min} px** (25th percentile of detected cell sizes)",
             pack_array(raw_image_np),
         )
 
@@ -744,17 +715,20 @@ def run_segmentation(image, model_choice, min_cell_size, max_cell_size,
             None,
             None,
             0.0,
-            gr.update(),
+            "",
             None,
         )
 
 
 def run_viability(stored_masks, stored_image_np):
-    """Run model-based viability classification. Returns overlay + counts + label_map."""
     if stored_masks is None or stored_image_np is None:
         return None, 0, 0, 0.0, "Please run segmentation first.", {}
+
+    # Lazy-load the classifier the first time viability is requested
+    ensure_viability_loaded()
     if VIABILITY_CLF is None:
-        return None, 0, 0, 0.0, "No viability model found. Add viability_clf.pkl and viability_scaler.pkl to the app directory.", {}
+        err = _VIABILITY_ERROR or "unknown error"
+        return None, 0, 0, 0.0, f"Viability model failed to load: {err}", {}
 
     masks    = unpack_array(stored_masks)
     image_np = unpack_array(stored_image_np)
@@ -773,14 +747,20 @@ def run_viability(stored_masks, stored_image_np):
 
 
 def pack_array(arr):
-    pil = Image.fromarray(arr.astype(np.uint8))
+    """
+    Serialise a numpy array to a base64 string for gr.State storage.
+    Uses numpy .npy format (preserves int32 exactly, no 255 truncation)
+    encoded as base64 so it is a plain Python string — safe for ZeroGPU
+    state serialisation which cannot handle raw bytes objects.
+    """
     buf = io.BytesIO()
-    pil.save(buf, format="PNG")
-    return buf.getvalue()
+    np.save(buf, arr)
+    return base64.b64encode(buf.getvalue()).decode("ascii")
 
 
 def unpack_array(data):
-    return np.array(Image.open(io.BytesIO(data)))
+    buf = io.BytesIO(base64.b64decode(data))
+    return np.load(buf, allow_pickle=False)
 
 
 def save_tab_result(cell_count, confluency, viab_percent):
@@ -820,29 +800,12 @@ def compute_summary(r1, r2, r3, r4):
     return avg_count, avg_conf, avg_viab, "\n".join(lines)
 
 
-# ---------------------------------------------------------------------------
+
 # Training data export — feature extraction per cell
-# ---------------------------------------------------------------------------
+
 
 def extract_cell_features(image_np, masks):
-    """
-    For every segmented cell, extract a fixed feature vector from the pixels
-    inside its mask.  Returns a list of dicts, one per cell.
-
-    Features:
-      RGB channels        — mean_r, mean_g, mean_b, std_r, std_g, std_b
-      HSV channels        — mean_h, mean_s, mean_v, std_s, std_v
-      Ratios              — blue_red_ratio, blue_green_ratio, rg_ratio
-      Morphology          — area_px, circularity
-      Centre/edge profile — inner_brightness, peak_brightness,
-                            bright_spot_fraction, ring_darkness,
-                            centre_periphery_ratio, brightness_std_normalised
-
-    Profile zones are tuned to hemocytometer live-cell morphology:
-    a small intense specular highlight at the centre surrounded by a dark
-    navy membrane ring. Dead cells are pale blue-grey blobs with no ring
-    and no bright spot.
-    """
+    
     if len(image_np.shape) == 2:
         image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2RGB)
     elif image_np.shape[2] == 4:
@@ -1003,9 +966,8 @@ def prepare_export(stored_masks, stored_image, threshold_bias):
     return path, msg
 
 
-# ---------------------------------------------------------------------------
+
 # Tab builder
-# ---------------------------------------------------------------------------
 
 def draw_polygon_overlay(image_pil, points):
     """
@@ -1063,14 +1025,12 @@ def clear_crop_points(image_pil):
 
 
 
-# ---------------------------------------------------------------------------
 # Label correction grid
-# ---------------------------------------------------------------------------
 
-THUMB_SIZE   = 80   # each cell thumbnail is THUMB_SIZE × THUMB_SIZE px
-GRID_COLS    = 8    # thumbnails per row
-BORDER       = 4    # coloured border thickness in px
-LABEL_H      = 16   # height of the text label strip at the bottom of each thumb
+THUMB_SIZE   = 80   
+GRID_COLS    = 10  
+BORDER       = 4    
+LABEL_H      = 16   
 
 def _crop_cell_thumb(image_np, masks, cid):
     """
@@ -1102,14 +1062,7 @@ def _crop_cell_thumb(image_np, masks, cid):
 
 
 def build_correction_grid(image_np, masks, labelled_features, raw_image_np=None):
-    """
-    Render all cell thumbnails into a single PIL image grid.
-    Each thumbnail has a coloured border: green=live(0), red=dead(1).
-    A small number in the corner identifies the cell_id.
-
-    Returns the PIL grid image.
-    Cell order in the grid matches the order of labelled_features.
-    """
+    
     if not labelled_features:
         placeholder = Image.fromarray(
             np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8)
@@ -1192,6 +1145,197 @@ def toggle_cell_label(labelled_features, image_np, masks, raw_image_np, evt: gr.
     return grid, updated, f"Tapped cell #{cell['cell_id']} → {'Dead' if cell['label']==1 else 'Live'}. {n_corrected} correction(s) total."
 
 
+# ---------------------------------------------------------------------------
+# Track A4 — self-correction review UI
+# ---------------------------------------------------------------------------
+
+SELF_CORRECTION_CLASS_COLORS = {
+    "split":  (255, 165, 0),   # orange
+    "merge":  (220, 50, 50),   # red
+    "debris": (150, 50, 220),  # purple
+}
+SC_LABEL_H = 30  # taller than the other grids' LABEL_H — two lines of text
+SELF_CORRECTION_REPAIR_TEXT = {
+    "removed": "removed",
+    "watershed_split": "split apart",
+    "merged_with_neighbor": "merged",
+    "already_resolved": "resolved w/ sibling",
+    "flagged_for_review": "needs review",
+}
+
+
+def draw_self_correction_overlay(base_image_np, masks_before, logs):
+    """Enumerates only the flagged (non-'correct') instances, colored by
+    predicted class and labeled with the model's calibrated confidence —
+    mirrors draw_viability_overlay's contour+label convention
+    (app_final.py:226) so a flagged cell can be matched back to its native
+    location. `base_image_np` is typically build_colored_mask_overlay's
+    output (the same "Segmentation Result" rendering), so the outlines draw
+    on top of the existing segmentation view rather than the raw photo.
+    `masks_before` is the pre-correction mask array (the crop location a
+    repaired/removed instance had before this run)."""
+    overlay = base_image_np.copy()
+    for log in logs:
+        color = SELF_CORRECTION_CLASS_COLORS.get(log.predicted_class, (200, 200, 200))
+        cell_mask = (masks_before == log.instance_id).astype(np.uint8)
+        contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(overlay, contours, -1, color, thickness=2)
+
+        ys, xs = np.where(cell_mask)
+        if len(ys) == 0:
+            continue
+        cx, cy = int(xs.mean()), int(ys.mean())
+        label_str = f"{log.predicted_class[0].upper()} {log.confidence:.0%}"
+        font, font_scale, thickness = cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1
+        (tw, th), _ = cv2.getTextSize(label_str, font, font_scale, thickness)
+        cv2.rectangle(overlay, (cx - tw // 2 - 1, cy - th // 2 - 1),
+                      (cx + tw // 2 + 1, cy + th // 2 + 1), (0, 0, 0), -1)
+        cv2.putText(overlay, label_str, (cx - tw // 2, cy + th // 2),
+                    font, font_scale, color, thickness, cv2.LINE_AA)
+    return overlay
+
+
+def build_self_correction_grid(image_np, masks_before, logs):
+    """Thumbnail grid of every flagged instance — mirrors build_correction_grid
+    (app_final.py:996) but reports the model's predicted class + calibrated
+    confidence + repair outcome (read-only review, not tap-to-toggle, since
+    the repair itself is auto-applied or explicitly left for review)."""
+    if not logs:
+        return Image.fromarray(np.zeros((THUMB_SIZE, THUMB_SIZE, 3), dtype=np.uint8))
+
+    n_cols = GRID_COLS
+    n_rows = (len(logs) + n_cols - 1) // n_cols
+    cell_w = THUMB_SIZE + 2 * BORDER
+    cell_h = THUMB_SIZE + 2 * BORDER + SC_LABEL_H
+
+    grid = Image.new("RGB", (n_cols * cell_w, n_rows * cell_h), (30, 30, 30))
+    draw = ImageDraw.Draw(grid)
+
+    for idx, log in enumerate(logs):
+        color = SELF_CORRECTION_CLASS_COLORS.get(log.predicted_class, (200, 200, 200))
+        thumb = _crop_cell_thumb(image_np, masks_before, log.instance_id)
+
+        col, row = idx % n_cols, idx // n_cols
+        x0, y0 = col * cell_w, row * cell_h
+
+        draw.rectangle([x0, y0, x0 + cell_w - 1, y0 + cell_h - 1], outline=color, width=BORDER)
+        grid.paste(thumb, (x0 + BORDER, y0 + BORDER))
+
+        strip_y = y0 + BORDER + THUMB_SIZE
+        draw.rectangle([x0, strip_y, x0 + cell_w - 1, y0 + cell_h - 1], fill=(20, 20, 20))
+        repaired = log.repair_applied not in ("flagged_for_review", "already_resolved")
+        status_mark = "✓" if repaired else "⚑"  # checkmark / flag
+        status_text = SELF_CORRECTION_REPAIR_TEXT.get(log.repair_applied, log.repair_applied)
+        draw.text((x0 + BORDER + 2, strip_y + 1),
+                  f"#{log.instance_id} {log.predicted_class[:1].upper()} {log.confidence:.0%}",
+                  fill=color)
+        draw.text((x0 + BORDER + 2, strip_y + 12),
+                  f"{status_mark} {status_text}",
+                  fill=color)
+
+    return grid
+
+
+def build_self_correction_details(logs):
+    """Full-text per-instance breakdown (class, confidence, outcome, and —
+    critically — *why* an instance was left for review) as a markdown table.
+    The thumbnail grid's text is too small to fit a real explanation; this is
+    the readable counterpart, directly answering 'why did/didn't this cell
+    get auto-fixed'."""
+    if not logs:
+        return ""
+    rows = ["| # | Class | Confidence | Outcome | Reason |", "|---|---|---|---|---|"]
+    for log in logs:
+        outcome = SELF_CORRECTION_REPAIR_TEXT.get(log.repair_applied, log.repair_applied)
+        reason = log.reason or "—"
+        rows.append(f"| {log.instance_id} | {log.predicted_class} | {log.confidence:.0%} | {outcome} | {reason} |")
+    return "\n".join(rows)
+
+
+def run_self_correction_step(use_correction, stored_masks, stored_image):
+    """Called after run_segmentation (opt-in, off by default per the guide's
+    Step 9). Classifies every instance and auto-applies the matching CPU
+    repair above its confidence threshold (error_detection/inference.py's
+    DEFAULT_THRESHOLDS: split/debris ~75%, merge kept elevated at 90% since
+    real merge examples are still scarce). Split repair additionally passes
+    through a geometric gate (flat-edge-fraction + combined-area check) that
+    tells a real over-segmentation split apart from two genuinely separate,
+    correctly-segmented cells that merely touch. Anything left unresolved is
+    reported for manual review rather than guessed at.
+
+    Updates the main 'Segmentation Result' image in place with colored
+    outlines on every flagged cell, and produces a second 'Corrected
+    Segmentation' panel rendering the masks after repair (splits merged back
+    into one color, debris removed, merges split into distinct colors) —
+    the before/after pair needed to actually judge repair quality."""
+    hidden_after = gr.update(value=None, visible=False)
+    if not use_correction or stored_masks is None or stored_image is None:
+        return (stored_masks, gr.update(), hidden_after, gr.update(value=None, visible=False),
+                gr.update(value="", visible=False), "", [], None,
+                gr.update(), gr.update(visible=False))
+
+    masks_before = unpack_array(stored_masks)
+    image_np = unpack_array(stored_image)
+
+    ensure_self_correction_loaded()
+    if SELF_CORRECTION_MODEL is None:
+        msg = f"Self-correction model unavailable ({_SELF_CORRECTION_ERROR}) — skipped."
+        return (stored_masks, gr.update(), hidden_after, gr.update(value=None, visible=False),
+                gr.update(value=msg, visible=True), "", [], None,
+                gr.update(), gr.update(visible=True))
+
+    corrected_masks, logs = SELF_CORRECTION_RUN_FN(
+        image_np, masks_before, SELF_CORRECTION_MODEL, temperature=SELF_CORRECTION_TEMP
+    )
+
+    n_repaired = sum(1 for log in logs if log.repair_applied not in ("flagged_for_review", "already_resolved"))
+    n_flagged = sum(1 for log in logs if log.repair_applied == "flagged_for_review")
+    if logs:
+        msg = (f"Self-correction: {len(logs)} instance(s) flagged — "
+               f"{n_repaired} auto-repaired, {n_flagged} left for manual review "
+               f"(split repairs are additionally checked against cell shape — "
+               f"see the table below for why any cell wasn't auto-fixed). "
+               f"'⚑ review' cells in the grid were not modified.")
+    else:
+        msg = "Self-correction: no split/merge/debris errors flagged above threshold."
+
+    before_overlay = build_colored_mask_overlay(image_np, masks_before)
+    annotated_before = draw_self_correction_overlay(before_overlay, masks_before, logs)
+    after_overlay = build_colored_mask_overlay(image_np, corrected_masks)
+    grid = build_self_correction_grid(image_np, masks_before, logs)
+    details_md = build_self_correction_details(logs)
+    new_cell_count = int(len(np.unique(corrected_masks)) - 1)
+
+    return (pack_array(corrected_masks),
+            gr.update(value=Image.fromarray(annotated_before)),
+            gr.update(value=Image.fromarray(after_overlay), visible=True),
+            gr.update(value=grid, visible=True),
+            gr.update(value=msg, visible=True),
+            details_md, logs, stored_masks,
+            new_cell_count, gr.update(visible=True))
+
+
+def revert_self_correction_step(pre_correction_masks, stored_image):
+    """Restores the pre-correction masks saved by run_self_correction_step,
+    and resets the Segmentation Result view + hides the correction panels."""
+    if pre_correction_masks is None:
+        return (pre_correction_masks, gr.update(), gr.update(value=None, visible=False),
+                gr.update(), gr.update(value=None, visible=False), gr.update(value="", visible=False),
+                gr.update(value="Nothing to revert.", visible=True))
+    masks = unpack_array(pre_correction_masks)
+    new_cell_count = int(len(np.unique(masks)) - 1)
+    plain_overlay = None
+    if stored_image is not None:
+        image_np = unpack_array(stored_image)
+        plain_overlay = Image.fromarray(build_colored_mask_overlay(image_np, masks))
+    return (pre_correction_masks, new_cell_count,
+            gr.update(value=None, visible=False),
+            gr.update(value=plain_overlay) if plain_overlay is not None else gr.update(),
+            gr.update(value=None, visible=False),
+            gr.update(value="", visible=False),
+            gr.update(value="Self-correction reverted — original segmentation restored.", visible=True))
+
+
 def prepare_export_corrected(stored_masks, stored_image, labelled_features, label_map):
     """Export CSV using labelled_features with any manual corrections applied."""
     if stored_masks is None or stored_image is None:
@@ -1256,14 +1400,29 @@ def build_tab(tab_index, masks_state, image_state, result_state):
                     value="Hemocytometer Model"
                 )
 
+                gr.Markdown("### Size Filters")
+
+                use_min_filter = gr.Checkbox(
+                    label="Enable minimum size filter",
+                    value=False,
+                    info="Remove objects smaller than the threshold below"
+                )
                 min_size_slider = gr.Slider(
                     minimum=0,
                     maximum=500,
                     value=0,
                     step=10,
-                    label="Minimum Cell Size (pixels). Leave at zero for automated recommendation",
+                    label="Minimum Cell Size (pixels)",
+                )
+                min_size_recommendation = gr.Markdown(
+                    value="*Run segmentation to see recommended minimum*",
                 )
 
+                use_max_filter = gr.Checkbox(
+                    label="Enable maximum size filter",
+                    value=False,
+                    info="Remove objects larger than the threshold below"
+                )
                 max_size_slider = gr.Slider(
                     minimum=0,
                     maximum=10000,
@@ -1311,6 +1470,20 @@ def build_tab(tab_index, masks_state, image_state, result_state):
                         info="Width of top exclusion zone"
                     )
 
+                gr.Markdown("### Automated Error Correction (Track A4 — experimental)")
+                use_self_correction = gr.Checkbox(
+                    label="Run self-correction after segmentation",
+                    value=False,
+                    info=(
+                        "Auto-repairs high-confidence split/merge/debris errors "
+                        "(CPU only, no extra GPU cost — see the review grid below "
+                        "for what changed). Merge repair uses a conservative "
+                        "confidence threshold since real merge examples are still "
+                        "scarce; anything below threshold is left flagged for you "
+                        "to review instead of being auto-applied."
+                    ),
+                )
+
                 segment_btn = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
 
             with gr.Column():
@@ -1319,6 +1492,30 @@ def build_tab(tab_index, masks_state, image_state, result_state):
                 overlay_out    = gr.Image(type="pil", label="Segmentation Result")
                 info_out       = gr.Textbox(label="Processing Info", lines=4)
 
+        with gr.Group(visible=False) as self_correction_section:
+            gr.Markdown(
+                "### Self-Correction Review\n"
+                "The **Segmentation Result** image above is outlined per flagged cell "
+                "(orange=split · red=merge · purple=debris) once self-correction runs. "
+                "Compare it against **Corrected Segmentation** below to judge repair quality — "
+                "splits should merge back into one color, merges should split into two, "
+                "debris should disappear."
+            )
+            self_correction_status = gr.Markdown(visible=False)
+            corrected_overlay_out = gr.Image(
+                type="pil", label="Corrected Segmentation (after auto-repair)", visible=False
+            )
+            self_correction_grid = gr.Image(
+                type="pil",
+                label="Flagged cells (orange=split · red=merge · purple=debris — ✓ fixed / ⚑ needs review)",
+                visible=False,
+            )
+            self_correction_details = gr.Markdown(visible=True)
+            revert_self_correction_btn = gr.Button("↩ Revert self-correction", size="sm")
+
+        correction_log_state       = gr.State(value=[])
+        pre_correction_masks_state = gr.State(value=None)
+
         with gr.Group(visible=False) as viability_section:
             gr.Markdown("### Viability Assessment (Trypan Blue)")
 
@@ -1405,9 +1602,27 @@ def build_tab(tab_index, masks_state, image_state, result_state):
         segment_btn.click(
             fn=run_segmentation,
             inputs=[img_input, model_dropdown, min_size_slider, max_size_slider,
+                    use_min_filter, use_max_filter,
                     use_stereo, left_excl, top_excl, crop_points_state],
             outputs=[cell_count_out, overlay_out, info_out, viability_section,
-                     masks_state, image_state, confluency_out, min_size_slider, raw_image_state]
+                     masks_state, image_state, confluency_out, min_size_recommendation, raw_image_state]
+        ).then(
+            # Opt-in step between segmentation and viability, per the guide's
+            # Step 9 — off unless use_self_correction is checked.
+            fn=run_self_correction_step,
+            inputs=[use_self_correction, masks_state, image_state],
+            outputs=[masks_state, overlay_out, corrected_overlay_out, self_correction_grid,
+                     self_correction_status, self_correction_details,
+                     correction_log_state, pre_correction_masks_state,
+                     cell_count_out, self_correction_section]
+        )
+
+        revert_self_correction_btn.click(
+            fn=revert_self_correction_step,
+            inputs=[pre_correction_masks_state, image_state],
+            outputs=[masks_state, cell_count_out, corrected_overlay_out,
+                     overlay_out, self_correction_grid, self_correction_details,
+                     self_correction_status]
         )
 
         # ---- Run Viability button -------------------------------------------
@@ -1496,9 +1711,8 @@ def build_tab(tab_index, masks_state, image_state, result_state):
 
 
 
-# ---------------------------------------------------------------------------
 # Gradio interface
-# ---------------------------------------------------------------------------
+
 with gr.Blocks(
     title="CellposeCellCounter",
     theme=gr.themes.Soft(),
@@ -1540,6 +1754,12 @@ with gr.Blocks(
             outputs=[avg_count_out, avg_conf_out, avg_viab_out, summary_box]
         )
 
+    # -------------------------------------------------------------------------
+    # Tab 6 — Error Labeling (QA) — builds a real hand-labeled validation set
+    # for error_detection/train.py --real-val-dir and calibrate.py --val-dir
+    # -------------------------------------------------------------------------
+    build_error_labeling_tab(MODEL_OPTIONS, run_segmentation, pack_array, unpack_array)
+
 
 
 if __name__ == "__main__":