Upload app.py

changed layout

app.py

@@ -1595,562 +1595,391 @@ def prepare_export_corrected(stored_masks, stored_image, labelled_features, labe
            f"{corrected} label(s) manually corrected.")
     return path, msg
 
-def build_tab(tab_index, masks_state, image_state, result_state):
-    with gr.Tab(f"Tab {tab_index}"):
-        gr.Markdown("Run segmentation")
-
-        # Per-tab state: list of (x,y) crop polygon points
-        crop_points_state = gr.State(value=[])
-        # Clean copy of the uploaded image (no polygon drawn on it)
-        base_image_state  = gr.State(value=None)
-        # preview->original scale, so taps on the small picker map back to
-        # full-resolution coordinates for the segmentation warp
-        preview_scale_state = gr.State(value=1.0)
-        # None = showing the overview; dict = showing a zoomed window
-        zoom_state = gr.State(value=None)
-        # segmentation wall time and the overlay frame, both needed by Tab 5
-        seg_time_state    = gr.State(value=0.0)
-        seg_overlay_state = gr.State(value=None)
-        #raw image state
-        raw_image_state  = gr.State(value=None)
-        # Masks as segmented, before stereological exclusion — lets the zones
-        # be re-chosen after segmentation without re-running the model.
-        base_masks_state = gr.State(value=None)
+# ===========================================================================
+# UI
+# ===========================================================================
+HEMO_MODEL    = "Hemocytometer Model"
+GENERAL_MODEL = "General Model"
 
-        with gr.Row():
-            with gr.Column():
-                img_input = gr.Image(
-                    type="pil",
-                    label="Upload image",
-                    image_mode="RGB",
-                    height=512
-                )
-
-                gr.Markdown(
-                    "### Crop region (optional)\n"
-                    "For each corner, tap **twice**: once roughly near it, then "
-                    "again precisely in the zoomed view that appears. The rough "
-                    "tap needs no accuracy — your finger can cover it entirely. "
-                    "Four corners define the region to segment; leave empty to "
-                    "use the whole image."
-                )
-
-                crop_display = gr.Image(
-                    type="pil",
-                    label="Tap to set crop vertices (up to 4)",
-                    interactive=True,
-                    height=400,
-                    format="jpeg",          # PNG costs ~3 s per redraw at 12 MP
-                    show_download_button=False,
-                )
-
-                crop_status = gr.Markdown("*Upload an image to enable cropping*")
-
-                clear_crop_btn = gr.Button("✕ Clear crop points", size="sm")
-
-                cancel_zoom_btn = gr.Button("↩ Back to full view", size="sm",
-                                            visible=False)
-
-                model_dropdown = gr.Dropdown(
-                    choices=list(MODEL_OPTIONS.keys()),
-                    label="Select Model",
-                    value="Hemocytometer Model"
-                )
-
-                gr.Markdown("### Size Filters")
-
-                use_min_filter = gr.Checkbox(
-                    label="Enable minimum size filter",
-                    value=True,
-                    info="Remove objects smaller than the threshold below. "
-                         "At 0 the app uses its own recommendation "
-                         "(25th percentile of detected object sizes)."
-                )
-                min_size_slider = gr.Slider(
-                    minimum=0,
-                    maximum=500,
-                    value=0,
-                    step=10,
-                    label="Minimum Cell Size (pixels)",
-                    interactive=True,
-                )
-                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,
-                    value=10000,
-                    step=10,
-                    label="Maximum Cell Size (pixels)",
-                    interactive=False,
-                )
-
-                segment_btn = gr.Button("🔬 Run Segmentation", variant="primary", size="lg")
-
-            with gr.Column():
-                cell_count_out = gr.Number(label="Total Cells Detected", precision=0)
-                confluency_out = gr.Number(label="Confluency (%)", precision=1)
-                overlay_out    = gr.Image(type="pil", label="Segmentation Result")
-                info_out       = gr.Textbox(label="Processing Info", lines=4)
-
-                # Applied after segmentation, so it lives beside the result it
-                # modifies rather than with the pre-segmentation settings.
-                gr.Markdown("### Stereological Counting")
-                use_stereo = gr.Checkbox(
-                    label="Enable Stereological Counting",
-                    value=True,
-                    info="Applied to the segmentation above — no re-segmenting needed"
-                )
-
-                with gr.Group(visible=True) as stereo_controls:
-                    gr.Markdown("""
-                    **Stereological Counting Rules:**
-                    - Cells touching LEFT or TOP exclusion zones are EXCLUDED
-                    - Cells touching RIGHT or BOTTOM edges are INCLUDED
-                    - This provides unbiased counting for quantification
-
-                    Zone widths are percentages of the **segmented image** —
-                    the cropped and perspective-corrected region, not the
-                    original upload. The red zones are drawn on the
-                    segmentation result above and update as you drag.
-                    """)
-
-                    left_excl = gr.Slider(
-                        minimum=0,
-                        maximum=50,
-                        value=1,
-                        step=1,
-                        label="Left Exclusion Width (%)",
-                        info="Width of left exclusion zone"
-                    )
-
-                    top_excl = gr.Slider(
-                        minimum=0,
-                        maximum=50,
-                        value=1,
-                        step=1,
-                        label="Top Exclusion Width (%)",
-                        info="Width of top exclusion zone"
-                    )
-
-        with gr.Group(visible=False) as viability_section:
-            gr.Markdown("### Viability Assessment (Trypan Blue)")
-
-            viab_run_btn   = gr.Button("Run Viability Analysis", variant="primary")
-
-            with gr.Row():
-                live_count_out = gr.Number(label="Live Cells (Green)", precision=0)
-                dead_count_out = gr.Number(label="Dead Cells (Red)",   precision=0)
-
-            viab_overlay     = gr.Image(type="pil", label="Viability (Green=Live · Red=Dead)")
-            viab_percent_out = gr.Number(label="Viability (%)", precision=1)
-            with gr.Row():
-                viab_info = gr.Textbox(label="Analysis Results", lines=5)
-                conc_1_1 = gr.Textbox(
-                    label=f"Concentration — {CONC_PRESETS[0][0]}",
-                    lines=5, interactive=False,
-                )
-                conc_1_10 = gr.Textbox(
-                    label=f"Concentration — {CONC_PRESETS[1][0]}",
-                    lines=5, interactive=False,
-                )
-
-            gr.Markdown("### Label Correction & Export")
-            gr.Markdown(
-                "After running viability, click **Build correction grid** to review every cell. "
-                "**Green border = Live, Red border = Dead** (model predictions). "
-                "Tap any thumbnail to flip its label — the counts and overlay update instantly. "
-                "Export the corrected CSV for retraining."
-            )
-
-            build_grid_btn    = gr.Button("🔲 Build correction grid", variant="secondary")
-            labelled_state    = gr.State(value=[])
-            label_map_state   = gr.State(value={})
-
-            correction_grid   = gr.Image(
-                type="pil",
-                label="Tap a cell to flip its label  (green=live · red=dead)",
-                interactive=True,
-                visible=False,
-            )
-            correction_status = gr.Markdown(visible=False)
-
-            with gr.Row():
-                export_btn  = gr.Button("⬇️ Export corrected CSV", variant="secondary")
-                export_info = gr.Textbox(label="Export status", lines=2, interactive=False)
-            export_file = gr.File(label="Download CSV", visible=False)
-
-        # ---- Event handlers ------------------------------------------------
-
-        def on_model_change(model_choice):
-            """Hemocytometer counting is only unbiased with the stereological
-            rules applied, so selecting that model turns them on. Left as a
-            normal checkbox rather than locked, so the setting can still be
-            switched off deliberately."""
-            if model_choice == "Hemocytometer Model":
-                return gr.update(value=True), gr.update(visible=True)
-            return gr.update(), gr.update()
-
-        model_dropdown.change(
-            fn=on_model_change,
-            inputs=[model_dropdown],
-            outputs=[use_stereo, stereo_controls]
-        )
 
-        use_stereo.change(
-            fn=toggle_stereological_mode,
-            inputs=[use_stereo],
-            outputs=[stereo_controls]
-        )
-
-        def on_image_upload(img):
-            if img is None:
-                return None, None, 1.0, "*Upload an image to enable cropping*"
-            preview, scale = make_preview(img)
-            return (preview, preview, scale,
-                    "*Image loaded — tap up to 4 points to define crop region*")
-
-        img_input.change(
-            fn=on_image_upload,
-            inputs=[img_input],
-            outputs=[crop_display, base_image_state, preview_scale_state, crop_status]
-        ).then(fn=lambda: ([], None), outputs=[crop_points_state, zoom_state])
-
-        # Grey the size sliders out while their filter is disabled, so the
-        # checkbox state is visible rather than silently inferred.
-        use_min_filter.change(
-            fn=lambda on: gr.update(interactive=bool(on)),
-            inputs=[use_min_filter], outputs=[min_size_slider]
-        )
-        use_max_filter.change(
-            fn=lambda on: gr.update(interactive=bool(on)),
-            inputs=[use_max_filter], outputs=[max_size_slider]
-        )
+def _crop_picker(label):
+    """Upload + tap-to-zoom corner picker. Returns the components and states.
 
-        def on_exclusion_change(stored_base, stored_image, use_stereo_on,
-                                left_pct, top_pct):
-            """Re-apply the exclusion zones to existing masks — no re-segmentation."""
-            if stored_base is None or stored_image is None:
-                return (gr.update(), gr.update(), gr.update(),
-                        gr.update(), gr.update())
-            base_masks = unpack_array(stored_base)
-            image_np = unpack_array(stored_image)
-            masks, cell_count, confluency, overlay_pil, excluded = render_segmentation(
-                base_masks, image_np, use_stereo_on, left_pct, top_pct
-            )
-            if use_stereo_on:
-                msg = (f"Stereological exclusion (Left: {left_pct}%, Top: {top_pct}%): "
-                       f"{excluded} cells excluded, {cell_count} counted.\n"
-                       f"Confluency: {confluency:.1f}%\n"
-                       f"Re-run viability classification to update live/dead counts.")
-            else:
-                msg = (f"Stereological exclusion off — {cell_count} cells counted.\n"
-                       f"Confluency: {confluency:.1f}%")
-            return cell_count, overlay_pil, confluency, pack_array(masks), msg
-
-        exclusion_inputs = [base_masks_state, image_state, use_stereo,
-                            left_excl, top_excl]
-        exclusion_outputs = [cell_count_out, overlay_out, confluency_out,
-                             masks_state, info_out]
-        for _component in (use_stereo, left_excl, top_excl):
-            _component.change(fn=on_exclusion_change,
-                inputs=exclusion_inputs, outputs=exclusion_outputs)
-
-        def _overview(preview, points, scale):
-            pv = [(int(x * scale), int(y * scale)) for x, y in points]
-            return draw_polygon_overlay(preview, pv)
-
-        def on_crop_click(full_img, preview, points, scale, zoom,
-                          evt: gr.SelectData):
-            """Two stages per corner: rough tap -> zoomed view -> precise tap."""
-            points = list(points or [])
-            if preview is None or full_img is None:
-                return (gr.update(), points, zoom, gr.update(), gr.update())
-
-            # Gradio can deliver an empty selection (gradio-app/gradio#5945)
-            if evt is None or evt.index is None or evt.index[0] is None:
-                return (gr.update(), points, zoom,
-                        "*Tap not registered — try again inside the image*",
-                        gr.update())
-
-            tx, ty = int(evt.index[0]), int(evt.index[1])
-
-            if zoom is None:
-                # ---- stage 1: rough tap on the overview ----------------------
-                if len(points) >= 4:
-                    return (_overview(preview, points, scale), points, None,
-                            "*4 points set ✓ — **✕ Clear** to redo, or run segmentation*",
-                            gr.update(visible=False))
-                view, z = make_zoom_view(full_img, tx / scale, ty / scale)
-                return (view, points, z,
-                        f"*Zoomed in — now tap corner {len(points) + 1} precisely "
-                        f"(the red cross is your rough tap)*",
-                        gr.update(visible=True))
-
-            # ---- stage 2: precise tap inside the zoomed view -----------------
-            x_full = zoom["x0"] + tx / zoom["zscale"]
-            y_full = zoom["y0"] + ty / zoom["zscale"]
-            W, H = full_img.size
-            new_points = points + [(int(min(max(0, x_full), W - 1)),
-                                    int(min(max(0, y_full), H - 1)))]
-            n = len(new_points)
-            status = (f"*{n} / 4 corners set — tap roughly near corner {n + 1}*"
-                      if n < 4 else
-                      "*4 points set ✓ — **✕ Clear** to redo, or run segmentation*")
-            return (_overview(preview, new_points, scale), new_points, None,
-                    status, gr.update(visible=False))
-
-        crop_display.select(fn=on_crop_click,
-            inputs=[img_input, base_image_state, crop_points_state,
-                    preview_scale_state, zoom_state],
-            outputs=[crop_display, crop_points_state, zoom_state, crop_status,
-                     cancel_zoom_btn])
-
-        def on_cancel_zoom(preview, points, scale):
-            n = len(points or [])
-            return (_overview(preview, points or [], scale), None,
-                    f"*Back to full view — {n} / 4 corners set*",
-                    gr.update(visible=False))
-
-        cancel_zoom_btn.click(fn=on_cancel_zoom,
-            inputs=[base_image_state, crop_points_state, preview_scale_state],
-            outputs=[crop_display, zoom_state, crop_status, cancel_zoom_btn])
-
-        def on_clear_crop(base_img):
-            img, pts = clear_crop_points(base_img)
-            return (img, pts, None, "*Points cleared — tap roughly near corner 1*",
-                    gr.update(visible=False))
-
-        clear_crop_btn.click(fn=on_clear_crop,
-            inputs=[base_image_state],
-            outputs=[crop_display, crop_points_state, zoom_state, crop_status,
-                     cancel_zoom_btn])
-
-
-        # ---- Viability ------------------------------------------------------
-        def on_run_viability(stored_masks, stored_image, seg_seconds=None):
-            overlay, alive, dead, viab_pct, info, label_map = run_viability(
-                stored_masks, stored_image)
-            if seg_seconds:
-                info = f"{info}\nSegmentation time: {float(seg_seconds):.2f} s"
-            return (overlay, alive, dead, viab_pct, info, label_map,
-                    format_concentration(alive, CONC_PRESETS[0][1]),
-                    format_concentration(alive, CONC_PRESETS[1][1]))
-
-        viab_outputs = [viab_overlay, live_count_out, dead_count_out,
-                        viab_percent_out, viab_info, label_map_state,
-                        conc_1_1, conc_1_10]
-
-        viab_run_btn.click(
-            fn=on_run_viability,
-            inputs=[masks_state, image_state, seg_time_state],
-            outputs=viab_outputs
-        ).then(
-            fn=save_tab_result,
-            inputs=[cell_count_out, confluency_out, viab_percent_out,
-                    live_count_out, dead_count_out, seg_time_state,
-                    raw_image_state, seg_overlay_state, viab_overlay],
-            outputs=[result_state]
-        )
-
-        segment_btn.click(
-            fn=run_segmentation,
-            inputs=[img_input, model_dropdown, min_size_slider, max_size_slider,
-                    use_stereo, left_excl, top_excl, crop_points_state,
-                    use_min_filter, use_max_filter],
-            outputs=[cell_count_out, overlay_out, info_out, viability_section,
-                     masks_state, image_state, confluency_out,
-                     min_size_recommendation, raw_image_state, base_masks_state,
-                     seg_time_state, seg_overlay_state]
-        ).then(
-            # Viability is wanted on essentially every run, so do it without a
-            # second button press. The button stays for re-running after the
-            # exclusion sliders or the label corrections change.
-            fn=on_run_viability,
-            inputs=[masks_state, image_state, seg_time_state],
-            outputs=viab_outputs
-        ).then(
-            fn=save_tab_result,
-            inputs=[cell_count_out, confluency_out, viab_percent_out,
-                    live_count_out, dead_count_out, seg_time_state,
-                    raw_image_state, seg_overlay_state, viab_overlay],
-            outputs=[result_state]
-        )
-
-        # ---- Build correction grid -----------------------------------------
-        def on_build_grid(stored_masks, stored_image, label_map, stored_raw_image):
-            if stored_masks is None or stored_image is None or not label_map:
-                return (gr.update(visible=False), [],
-                        gr.update(value="*Run viability analysis first.*", visible=True))
-            masks        = unpack_array(stored_masks)
-            image_np     = unpack_array(stored_image)
-            raw_image_np = unpack_array(stored_raw_image) if stored_raw_image is not None else None
-            features     = extract_cell_features(image_np, masks)
-            labelled     = attach_viability_labels(features, masks, image_np, label_map)
-            if not labelled:
-                return (gr.update(visible=False), [],
-                        gr.update(value="*No cells found.*", visible=True))
-            grid = build_correction_grid(image_np, masks, labelled, raw_image_np)
-            n    = len(labelled)
-            dead = sum(1 for r in labelled if r["label"] == 1)
-            msg  = (f"*{n} cells — {n-dead} live (green), {dead} dead (red). "
-                    f"Tap any thumbnail to flip its label.*")
-            return gr.update(value=grid, visible=True), labelled, gr.update(value=msg, visible=True)
-
-        build_grid_btn.click(
-            fn=on_build_grid,
-            inputs=[masks_state, image_state, label_map_state, raw_image_state],
-            outputs=[correction_grid, labelled_state, correction_status]
-        )
-
-        # ---- Grid tap — flip label, update overlay + counts ----------------
-        def on_grid_tap(labelled, stored_masks, stored_image, stored_raw_image, evt: gr.SelectData):
-            if not labelled or stored_masks is None:
-                return None, labelled, "", 0, 0, 0.0, None, {}
-            masks        = unpack_array(stored_masks)
-            image_np     = unpack_array(stored_image)
-            raw_image_np = unpack_array(stored_raw_image) if stored_raw_image is not None else None
-            grid, updated, msg = toggle_cell_label(labelled, image_np, masks, raw_image_np, evt)
-
-            # Rebuild label_map from corrected labelled list
-            new_label_map = {int(f["cell_id"]): int(f["label"]) for f in updated}
-            overlay_np    = draw_viability_overlay(image_np, masks, new_label_map)
-            dead  = sum(1 for f in updated if f["label"] == 1)
-            alive = len(updated) - dead
-            total = alive + dead
-            viab_pct = (alive / total * 100) if total > 0 else 0.0
-
-            return (grid, updated, f"*{msg}*",
-                    alive, dead, viab_pct,
-                    Image.fromarray(overlay_np), new_label_map)
-
-        correction_grid.select(
-            fn=on_grid_tap,
-            inputs=[labelled_state, masks_state, image_state, raw_image_state],
-            outputs=[correction_grid, labelled_state, correction_status,
-                     live_count_out, dead_count_out, viab_percent_out,
-                     viab_overlay, label_map_state]
-        )
-
-        # ---- Export --------------------------------------------------------
-        def on_export(stored_masks, stored_image, labelled, label_map):
-            path, msg = prepare_export_corrected(stored_masks, stored_image, labelled, label_map)
-            if path is None:
-                return gr.update(visible=False), msg
-            return gr.update(value=path, visible=True), msg
-
-        export_btn.click(
-            fn=on_export,
-            inputs=[masks_state, image_state, labelled_state, label_map_state],
-            outputs=[export_file, export_info]
-        )
+    Factored out because Tab 1 needs four of these and Tab 2 needs one; the
+    wiring is identical, only the label differs.
+    """
+    img_input = gr.Image(type="pil", label=label, image_mode="RGB", height=220)
+    crop_display = gr.Image(
+        type="pil", label="Tap roughly, then precisely in the zoom",
+        interactive=True, height=340, format="jpeg",
+        show_download_button=False,
+    )
+    crop_status = gr.Markdown("*Upload an image to set the counting square*")
+    with gr.Row():
+        clear_btn  = gr.Button("✕ Clear", size="sm")
+        cancel_btn = gr.Button("↩ Back to full view", size="sm", visible=False)
+
+    st = {
+        "img": img_input,
+        "display": crop_display,
+        "status": crop_status,
+        "points": gr.State(value=[]),
+        "base": gr.State(value=None),
+        "scale": gr.State(value=1.0),
+        "zoom": gr.State(value=None),
+    }
 
+    def on_upload(img):
+        if img is None:
+            return None, None, 1.0, "*Upload an image to set the counting square*"
+        preview, scale = make_preview(img)
+        return preview, preview, scale, "*Tap roughly near corner 1*"
+
+    img_input.change(
+        fn=on_upload, inputs=[img_input],
+        outputs=[crop_display, st["base"], st["scale"], crop_status]
+    ).then(fn=lambda: ([], None), outputs=[st["points"], st["zoom"]])
+
+    def _overview(preview, points, scale):
+        return draw_polygon_overlay(
+            preview, [(int(x * scale), int(y * scale)) for x, y in points])
+
+    def on_click(full_img, preview, points, scale, zoom, evt: gr.SelectData):
+        points = list(points or [])
+        if preview is None or full_img is None:
+            return gr.update(), points, zoom, gr.update(), gr.update()
+        if evt is None or evt.index is None or evt.index[0] is None:
+            return (gr.update(), points, zoom,
+                    "*Tap not registered — try again inside the image*", gr.update())
+        tx, ty = int(evt.index[0]), int(evt.index[1])
+
+        if zoom is None:
+            if len(points) >= 4:
+                return (_overview(preview, points, scale), points, None,
+                        "*4 corners set ✓ — ✕ Clear to redo*",
+                        gr.update(visible=False))
+            view, z = make_zoom_view(full_img, tx / scale, ty / scale)
+            return (view, points, z,
+                    f"*Zoomed — tap corner {len(points) + 1} precisely*",
+                    gr.update(visible=True))
+
+        x = zoom["x0"] + tx / zoom["zscale"]
+        y = zoom["y0"] + ty / zoom["zscale"]
+        W, H = full_img.size
+        pts = points + [(int(min(max(0, x), W - 1)), int(min(max(0, y), H - 1)))]
+        n = len(pts)
+        msg = (f"*{n} / 4 corners — tap roughly near corner {n + 1}*" if n < 4
+               else "*4 corners set ✓*")
+        return (_overview(preview, pts, scale), pts, None, msg,
+                gr.update(visible=False))
+
+    crop_display.select(
+        fn=on_click,
+        inputs=[img_input, st["base"], st["points"], st["scale"], st["zoom"]],
+        outputs=[crop_display, st["points"], st["zoom"], crop_status, cancel_btn])
+
+    cancel_btn.click(
+        fn=lambda preview, pts, sc: (
+            _overview(preview, pts or [], sc), None,
+            f"*Back to full view — {len(pts or [])} / 4 corners*",
+            gr.update(visible=False)),
+        inputs=[st["base"], st["points"], st["scale"]],
+        outputs=[crop_display, st["zoom"], crop_status, cancel_btn])
+
+    clear_btn.click(
+        fn=lambda base: (base, [], None, "*Cleared — tap roughly near corner 1*",
+                         gr.update(visible=False)),
+        inputs=[st["base"]],
+        outputs=[crop_display, st["points"], st["zoom"], crop_status, cancel_btn])
+
+    return st
+
+
+def _segment_one(image, points, model_name, use_min, min_px,
+                 use_stereo, left_pct, top_pct, want_viability):
+    """Segment (and optionally classify) a single image. Returns a dict."""
+    if image is None:
+        return None
+    out = run_segmentation(image, model_name, min_px, 10000,
+                           use_stereo, left_pct, top_pct, points,
+                           use_min, False)
+    (cell_count, seg_overlay, info_msg, _vis, packed_masks, packed_img,
+     confluency, _rec, packed_raw, _base, seg_seconds, _ov) = out
+    if packed_masks is None:
+        return {"error": info_msg}
+
+    res = {"count": cell_count, "confluency": confluency,
+           "seg_overlay": seg_overlay, "seconds": seg_seconds,
+           "info": info_msg, "raw": packed_raw,
+           "masks": packed_masks, "img": packed_img, "label_map": None,
+           "alive": None, "dead": None, "viab": None, "viab_overlay": None}
+
+    if want_viability:
+        v_overlay, alive, dead, viab_pct, _vinfo, lm = run_viability(
+            packed_masks, packed_img)
+        res.update(alive=alive, dead=dead, viab=viab_pct,
+                   viab_overlay=v_overlay, label_map=lm)
+    return res
+
+
+def _build_correction(res):
+    """Per-cell thumbnails for manual live/dead correction.
+
+    Only called when the user has ticked the correction checkbox: the grid costs
+    one crop per cell, so building it for four squares that nobody inspects is
+    pure waste.
+    """
+    try:
+        masks = unpack_array(res["masks"])
+        img   = unpack_array(res["img"])
+        raw   = unpack_array(res["raw"]) if res.get("raw") is not None else None
+        feats = extract_cell_features(img, masks)
+        labelled = attach_viability_labels(feats, masks, img, res.get("label_map"))
+        return build_correction_grid(img, masks, labelled, raw), labelled
+    except Exception as e:
+        print("correction grid failed:", e)
+        return None, []
 
 
-# ---------------------------------------------------------------------------
-# Gradio interface
-# ---------------------------------------------------------------------------
-with gr.Blocks(
-    title="CellposeCellCounter",
-    theme=gr.themes.Soft(),
-) as demo:
+with gr.Blocks(title="CellposeCellCounter", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# CellposeCellCounter")
-    gr.Markdown("For accurate cell confluency, crop the image to display only desired area. Note that some image file types are not yet supported. PNG and JPEG are preferred.")
-
-    # Shared mask/image state (one pair per tab so tabs don't clobber each other)
-    masks_states  = [gr.State(value=None) for _ in range(4)]
-    image_states  = [gr.State(value=None) for _ in range(4)]
-    result_states = [gr.State(value=None) for _ in range(4)]
-
-    # Build Tabs 1–4 with a loop
-    for i in range(4):
-        build_tab(i + 1, masks_states[i], image_states[i], result_states[i])
-
-    # -------------------------------------------------------------------------
-    # Tab 5 — Summary
-    # -------------------------------------------------------------------------
-    with gr.Tab("Tab 5 — Summary"):
-        gr.Markdown("## Average Results Across All Tabs")
+
+    # =======================================================================
+    # Tab 1 — hemocytometer counting (4 squares, one run)
+    # =======================================================================
+    with gr.Tab("tab1_hemacyto_counting"):
         gr.Markdown(
-            "Run segmentation in one or more tabs, "
-            "then click **Refresh Summary** to see the averages."
+            "### Hemocytometer counting\n"
+            "Upload the **four corner squares**, set four corners on each, then "
+            "press **Run segmentation** once. Uses the hemocytometer model and "
+            "runs viability automatically."
         )
 
-        refresh_btn = gr.Button("🔄 Refresh Summary", variant="primary", size="lg")
-
+        with gr.Accordion("Settings", open=False):
+            h_use_min = gr.Checkbox(
+                label="Enable minimum size filter", value=True,
+                info="At 0 the app uses the 25th percentile of each image — "
+                     "set a fixed value for reproducible counts.")
+            h_min_px = gr.Slider(0, 500, value=0, step=10,
+                                 label="Minimum Cell Size (pixels)")
+            h_use_stereo = gr.Checkbox(
+                label="Enable stereological counting", value=True,
+                info="Cells touching the left/top exclusion lines are excluded.")
+            h_left = gr.Slider(0, 50, value=1, step=1, label="Left exclusion (%)")
+            h_top  = gr.Slider(0, 50, value=1, step=1, label="Top exclusion (%)")
+            h_correct = gr.Checkbox(
+                label="Enable manual label correction", value=False,
+                info="Adds a tap-to-flip thumbnail grid per square. Off by "
+                     "default because building it costs one crop per cell.")
+
+        pickers, seg_imgs, viab_imgs, stat_boxes = [], [], [], []
+        corr_accordions, corr_grids, corr_btns, corr_infos, corr_files = [], [], [], [], []
+        masks_sts, img_sts, raw_sts, lab_sts, lmap_sts = [], [], [], [], []
+        for _i in range(4):
+            with gr.Accordion(f"Square {_i + 1}", open=(_i == 0)):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        pickers.append(_crop_picker(f"Square {_i + 1} image"))
+                    with gr.Column(scale=1):
+                        seg_imgs.append(gr.Image(type="pil", label="Segmentation",
+                                                 height=260))
+                        viab_imgs.append(gr.Image(type="pil",
+                                                  label="Viability (green=live · red=dead)",
+                                                  height=260))
+                        stat_boxes.append(gr.Textbox(label="Results", lines=7,
+                                                     interactive=False))
+                with gr.Accordion("Label correction", open=False,
+                                  visible=False) as _acc:
+                    corr_accordions.append(_acc)
+                    gr.Markdown("Tap a thumbnail to flip live ↔ dead "
+                                "(green = live, red = dead).")
+                    corr_grids.append(gr.Image(
+                        type="pil", label="Correction grid",
+                        interactive=True, height=420, format="jpeg",
+                        show_download_button=False))
+                    with gr.Row():
+                        corr_btns.append(gr.Button("⬇️ Export corrected CSV",
+                                                   size="sm"))
+                        corr_infos.append(gr.Textbox(label="Export status",
+                                                     lines=2, interactive=False))
+                    corr_files.append(gr.File(label="Download CSV", visible=False))
+                masks_sts.append(gr.State(value=None))
+                img_sts.append(gr.State(value=None))
+                raw_sts.append(gr.State(value=None))
+                lab_sts.append(gr.State(value=None))
+                lmap_sts.append(gr.State(value=None))
+
+        run_all_btn = gr.Button("🔬 Run segmentation (all 4 squares)",
+                                variant="primary", size="lg")
+        run_status  = gr.Textbox(label="Progress", lines=2, interactive=False)
+
+        gr.Markdown("## Summary across the four squares")
         with gr.Row():
-            avg_count_out = gr.Number(label="Avg Cell Count",      precision=1)
-            avg_conf_out  = gr.Number(label="Avg Confluency (%)",  precision=1)
-            avg_viab_out  = gr.Number(label="Avg Viability (%)",   precision=1)
-
-        summary_box = gr.Textbox(label="Per-Tab Breakdown", lines=11)
-
-        gr.Markdown("### Cell Concentration (from the mean across tabs)")
+            s_count = gr.Number(label="Mean total cells", precision=1)
+            s_conf  = gr.Number(label="Mean confluency (%)", precision=1)
+            s_viab  = gr.Number(label="Mean viability (%)", precision=1)
+        s_break = gr.Textbox(label="Per-square breakdown", lines=11,
+                             interactive=False)
         with gr.Row():
-            summary_conc_2x = gr.Textbox(
-                label=f"Concentration — {CONC_PRESETS[0][0]}",
-                lines=8, interactive=False,
-            )
-            summary_conc_10x = gr.Textbox(
-                label=f"Concentration — {CONC_PRESETS[1][0]}",
-                lines=8, interactive=False,
-            )
-
-        refresh_btn.click(
-            fn=compute_summary,
-            inputs=result_states,   # list of 4 gr.State components
-            outputs=[avg_count_out, avg_conf_out, avg_viab_out, summary_box,
-                     summary_conc_2x, summary_conc_10x]
-        )
+            s_c2  = gr.Textbox(label=f"Concentration — {CONC_PRESETS[0][0]}",
+                               lines=8, interactive=False)
+            s_c10 = gr.Textbox(label=f"Concentration — {CONC_PRESETS[1][0]}",
+                               lines=8, interactive=False)
 
         gr.Markdown("### Export")
         with gr.Row():
-            summary_export_btn  = gr.Button("⬇️ Export summary CSV",
-                                            variant="secondary")
-            summary_export_info = gr.Textbox(label="Export status", lines=2,
-                                             interactive=False)
-        summary_export_file = gr.File(label="Download summary CSV", visible=False)
-
-        def on_export_summary(*tab_results):
-            path, msg = export_summary_csv(*tab_results)
-            if path is None:
-                return gr.update(visible=False), msg
-            return gr.update(value=path, visible=True), msg
-
-        summary_export_btn.click(
-            fn=on_export_summary,
-            inputs=result_states,
-            outputs=[summary_export_file, summary_export_info]
+            s_csv_btn = gr.Button("⬇️ Export summary CSV", variant="secondary")
+            s_pdf_btn = gr.Button("📄 Export summary PDF (with images)",
+                                  variant="secondary")
+        s_export_info = gr.Textbox(label="Export status", lines=2,
+                                   interactive=False)
+        s_csv_file = gr.File(label="Download CSV", visible=False)
+        s_pdf_file = gr.File(label="Download PDF", visible=False)
+
+        result_states = [gr.State(value=None) for _ in range(4)]
+
+        def run_all(i1, p1, i2, p2, i3, p3, i4, p4,
+                    use_min, min_px, use_stereo, left_pct, top_pct, want_correct):
+            imgs = [(i1, p1), (i2, p2), (i3, p3), (i4, p4)]
+            outs, sts, results, done, notes = [], [], [], 0, []
+
+            def _blank(msg):
+                outs.extend([None, None, msg, None])
+                sts.extend([None, None, None, None, None])
+                results.append(None)
+
+            for idx, (img, pts) in enumerate(imgs, start=1):
+                r = _segment_one(img, pts, HEMO_MODEL, use_min, min_px,
+                                 use_stereo, left_pct, top_pct, True)
+                if r is None:
+                    _blank("No image uploaded.")
+                    continue
+                if "error" in r:
+                    _blank(r["error"])
+                    notes.append(f"Square {idx} failed")
+                    continue
+                if not (pts and len(pts) >= 3):
+                    notes.append(f"Square {idx}: no crop set — whole image used")
+
+                grid, labelled = (_build_correction(r) if want_correct
+                                  else (None, None))
+
+                txt = (f"Total cells: {r['count']}\n"
+                       f"Live (green): {r['alive']}\n"
+                       f"Dead (red): {r['dead']}\n"
+                       f"Viability: {(r['viab'] or 0):.1f}%\n"
+                       f"Confluency: {r['confluency']:.1f}%\n"
+                       f"Segmentation time: {r['seconds']:.2f} s")
+                outs.extend([r["seg_overlay"], r["viab_overlay"], txt, grid])
+                sts.extend([r["masks"], r["img"], r["raw"], labelled,
+                            r["label_map"]])
+                results.append(save_tab_result(
+                    r["count"], r["confluency"], r["viab"], r["alive"], r["dead"],
+                    r["seconds"], r["raw"], r["seg_overlay"], r["viab_overlay"]))
+                done += 1
+
+            avg_c, avg_f, avg_v, brk, c2, c10 = compute_summary(*results)
+            status = f"Segmented {done} of 4 squares."
+            if notes:
+                status += " ⚠ " + "; ".join(notes)
+            return outs + sts + results + [avg_c, avg_f, avg_v, brk, c2, c10, status]
+
+        run_all_outputs = []
+        for _i in range(4):
+            run_all_outputs += [seg_imgs[_i], viab_imgs[_i], stat_boxes[_i],
+                                corr_grids[_i]]
+        for _i in range(4):
+            run_all_outputs += [masks_sts[_i], img_sts[_i], raw_sts[_i],
+                                lab_sts[_i], lmap_sts[_i]]
+        run_all_outputs += result_states
+        run_all_outputs += [s_count, s_conf, s_viab, s_break, s_c2, s_c10,
+                            run_status]
+
+        h_correct.change(
+            fn=lambda on: [gr.update(visible=bool(on))] * 4,
+            inputs=[h_correct], outputs=corr_accordions)
+
+        for _i in range(4):
+            corr_grids[_i].select(
+                fn=toggle_cell_label,
+                inputs=[lab_sts[_i], img_sts[_i], masks_sts[_i], raw_sts[_i]],
+                outputs=[corr_grids[_i], lab_sts[_i]])
+
+            def _make_export(k):
+                def _export(masks, img, labelled, lmap):
+                    path, msg = prepare_export_corrected(masks, img, labelled, lmap)
+                    return ((gr.update(visible=False) if path is None
+                             else gr.update(value=path, visible=True)), msg)
+                return _export
+
+            corr_btns[_i].click(
+                fn=_make_export(_i),
+                inputs=[masks_sts[_i], img_sts[_i], lab_sts[_i], lmap_sts[_i]],
+                outputs=[corr_files[_i], corr_infos[_i]])
+
+        run_all_btn.click(
+            fn=run_all,
+            inputs=[pickers[0]["img"], pickers[0]["points"],
+                    pickers[1]["img"], pickers[1]["points"],
+                    pickers[2]["img"], pickers[2]["points"],
+                    pickers[3]["img"], pickers[3]["points"],
+                    h_use_min, h_min_px, h_use_stereo, h_left, h_top,
+                    h_correct],
+            outputs=run_all_outputs)
+
+        def _csv(*rs):
+            path, msg = export_summary_csv(*rs)
+            return (gr.update(visible=False) if path is None
+                    else gr.update(value=path, visible=True)), msg
+
+        def _pdf(*rs):
+            path, msg = export_summary_pdf(*rs)
+            return (gr.update(visible=False) if path is None
+                    else gr.update(value=path, visible=True)), msg
+
+        s_csv_btn.click(fn=_csv, inputs=result_states,
+                        outputs=[s_csv_file, s_export_info])
+        s_pdf_btn.click(fn=_pdf, inputs=result_states,
+                        outputs=[s_pdf_file, s_export_info])
+
+    # =======================================================================
+    # Tab 2 — confluency only (general model, no viability)
+    # =======================================================================
+    with gr.Tab("tab2_confluency"):
+        gr.Markdown(
+            "### Confluency estimation\n"
+            "Uses the **general** model. Cropping is optional — leave the corners "
+            "unset to measure the whole field. No viability classification."
         )
-
         with gr.Row():
-            summary_pdf_btn  = gr.Button("📄 Export summary PDF (with images)",
-                                         variant="secondary")
-            summary_pdf_info = gr.Textbox(label="PDF status", lines=2,
-                                          interactive=False)
-        summary_pdf_file = gr.File(label="Download summary PDF", visible=False)
-
-        def on_export_summary_pdf(*tab_results):
-            path, msg = export_summary_pdf(*tab_results)
-            if path is None:
-                return gr.update(visible=False), msg
-            return gr.update(value=path, visible=True), msg
-
-        summary_pdf_btn.click(
-            fn=on_export_summary_pdf,
-            inputs=result_states,
-            outputs=[summary_pdf_file, summary_pdf_info]
-        )
+            with gr.Column(scale=1):
+                c_pick = _crop_picker("Image")
+                c_run  = gr.Button("🔬 Run segmentation", variant="primary",
+                                   size="lg")
+            with gr.Column(scale=1):
+                c_overlay = gr.Image(type="pil", label="Segmentation", height=340)
+                with gr.Row():
+                    c_count = gr.Number(label="Cells detected", precision=0)
+                    c_conf  = gr.Number(label="Confluency (%)", precision=2)
+                c_info = gr.Textbox(label="Processing info", lines=6,
+                                    interactive=False)
+
+        def run_confluency(img, pts):
+            if img is None:
+                return None, 0, 0.0, "Upload an image first."
+            r = _segment_one(img, pts, GENERAL_MODEL, False, 0,
+                             False, 0, 0, False)
+            if r is None or "error" in r:
+                return None, 0, 0.0, (r or {}).get("error", "Segmentation failed.")
+            note = "" if (pts and len(pts) >= 3) else "\nWhole image used (no crop set)."
+            return (r["seg_overlay"], r["count"], round(r["confluency"], 2),
+                    r["info"] + note)
 
+        c_run.click(fn=run_confluency,
+                    inputs=[c_pick["img"], c_pick["points"]],
+                    outputs=[c_overlay, c_count, c_conf, c_info])
 
 
 if __name__ == "__main__":
     _prefetch_models()
-    demo.launch()
+    demo.launch()