Sync from GitHub via hub-sync

app.py

@@ -14,11 +14,18 @@ import csv
 import joblib
 import os
 
-HF_REPO_ID = "myang4218/cellposemodel"
-HF_REPO_ID2 = "LiangLabUMB/viability_model"
+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 = {}
@@ -35,16 +42,13 @@ try:
 except Exception as e:
     print(f"Viability classifier not found or failed to load: {e}")
 
-# ---- mobile-safe size limits (aggressive for Safari) ----
+# 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 +156,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]
@@ -188,11 +188,7 @@ def classify_cells_by_model(image_np, masks):
 
 
 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 +216,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 +343,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 +355,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
@@ -567,7 +456,7 @@ def _segment_patch(args):
         model = models.CellposeModel(gpu=True, pretrained_model=model_path)
         loaded_models[model_filename] = model
 
-    mask, _, _ = model.eval(patch_np, diameter=None, channels=[0, 0])
+    mask, _, _ = model.eval(patch_np, diameter=None)
     return mask, row_start, col_start
 
 
@@ -579,7 +468,8 @@ def run_segmentation_patched(image_np, model_filename):
     that patching adds overhead without benefit.
     """
     h, w = image_np.shape[:2]
-    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=model_filename)
+    repo       = MODEL_REPOS.get(model_filename, HF_REPO_ID)
+    model_path = hf_hub_download(repo_id=repo, filename=model_filename)
     if model_filename in loaded_models:
         model = loaded_models[model_filename]
     else:
@@ -588,15 +478,16 @@ def run_segmentation_patched(image_np, model_filename):
 
     # Small images: no benefit from patching
     if max(h, w) <= MIN_PATCH_DIM * 2:
-        mask, _, _ = model.eval(image_np, diameter=None, channels=[0, 0])
+        mask, _, _ = model.eval(image_np, diameter=None)
         return mask, 1   # 1 patch
 
     patches = _split_patches(image_np)
     n_patches = len(patches)
 
     # Build argument list for the thread pool
+    patch_repo = MODEL_REPOS.get(model_filename, HF_REPO_ID)
     args_list = [
-        (patch, r, c, model_filename, HF_REPO_ID)
+        (patch, r, c, model_filename, patch_repo)
         for patch, r, c in patches
     ]
 
@@ -618,6 +509,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 +553,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 +576,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:
@@ -729,7 +618,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,13 +633,12 @@ 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.", {}
     if VIABILITY_CLF is None:
@@ -773,14 +661,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 bytes for gr.State storage.
+    Uses numpy's .npy format (not PNG) so integer dtypes of any
+    magnitude are preserved exactly — PNG is 8-bit only and silently
+    truncates cell IDs above 255.
+    """
     buf = io.BytesIO()
-    pil.save(buf, format="PNG")
+    np.save(buf, arr)
     return buf.getvalue()
 
 
 def unpack_array(data):
-    return np.array(Image.open(io.BytesIO(data)))
+    buf = io.BytesIO(data)
+    return np.load(buf, allow_pickle=False)
 
 
 def save_tab_result(cell_count, confluency, viab_percent):
@@ -820,29 +714,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 +880,8 @@ def prepare_export(stored_masks, stored_image, threshold_bias):
     return path, msg
 
 
-# ---------------------------------------------------------------------------
+
 # Tab builder
-# ---------------------------------------------------------------------------
 
 def draw_polygon_overlay(image_pil, points):
     """
@@ -1063,14 +939,12 @@ def clear_crop_points(image_pil):
 
 
 
-# ---------------------------------------------------------------------------
 # Label correction grid
-# ---------------------------------------------------------------------------
 
-THUMB_SIZE   = 80   # each cell thumbnail is THUMB_SIZE × THUMB_SIZE px
-GRID_COLS    = 6    # 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 +976,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)
@@ -1256,14 +1123,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,
@@ -1405,9 +1287,10 @@ 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]
         )
 
         # ---- Run Viability button -------------------------------------------
@@ -1496,9 +1379,8 @@ def build_tab(tab_index, masks_state, image_state, result_state):
 
 
 
-# ---------------------------------------------------------------------------
 # Gradio interface
-# ---------------------------------------------------------------------------
+
 with gr.Blocks(
     title="CellposeCellCounter",
     theme=gr.themes.Soft(),