Add edema and swim bladder segmentation and align the model presets with the webapp

CORE_SOURCE.md

@@ -0,0 +1,76 @@
+# Core source provenance
+
+`ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerCore/` is a manually-maintained
+port of the analysis engine from the reference webapp
+(`markdanielarndt/Zebrafish_webapp`, live deployment at
+`https://huggingface.co/spaces/markdanielarndt/Zebrafish`). This file exists
+so a future manual re-sync (webapp adds/changes a feature, this extension
+needs to catch up, or vice versa) is a targeted diff against a known state
+instead of a full re-read of both codebases.
+
+**Ground truth for "current webapp behavior" is always the live HF Space**
+(`.../raw/main/<file>.py`), not any local clone of the webapp repo — local
+clones have been observed to drift out of date (e.g. missing swim bladder
+segmentation, the edema UI, and the brush-based mask editor that are live in
+the deployed app as of 2026-07-20).
+
+## Provenance table
+
+| Core file (this repo) | Webapp source | Ported functions | State at last sync | Deliberate deviations |
+|---|---|---|---|---|
+| `seg.py` | `seg.py` | `_load_unet_model`, `segmentation_pipeline` | 2026-07-29 (webapp `app.py`/`seg.py`; the live HF Space and the webapp's GitHub repo were byte-identical for both files at that date) | `model_type` support (`"Unet"`/`"FPN"`) and the `include_swimbladder` path were re-synced on 2026-07-29; the pipeline's own defaults are kept identical to the webapp's so a future diff stays short. Return arity differs: the webapp enumerates the flag combinations explicitly, this port appends the requested masks in a fixed order (eyes, edema, swim bladder), which the caller unpacks by the flags it passed. No Hugging Face Hub download in this layer — `_load_unet_model` only accepts a local `model_path`/`filename`; `repo_id`/`revision`/`force_download` params are accepted for call-site compatibility but ignored. Downloading is `ZebrafishEmbryoAnalyzerLib/model_downloader.py`'s job, driven by `ZebrafishEmbryoAnalyzerLib/model_manifest.py` (webapp downloads directly via `huggingface_hub.hf_hub_download` inline). |
+| `seg_helper.py` | `seg.py` (helpers were inline in the webapp's single seg module at the time of the original port) | `load_images_from_path`, `segment_fish`, `fill_holes`, `grow_mask` | 2026-07-20 | None known — pure numpy/opencv, kept close to the original logic. |
+| `length.py` | `length.py` | `compute_eye_metrics`, `compute_eye_diameters`, `tube_length_border2border`, `classification_curvature`, `load_model`, `preprocess_masked_image`, `compute_tube_metrics` | 2026-07-29 (`compute_tube_metrics`; rest 2026-07-20) | `compute_tube_metrics` ported near-verbatim on 2026-07-29 from the live Space; only deviation is a deferred `import cv2` inside the function body, matching this repo's lazy-import convention for compiled extensions. `select_torch_device` is ours, not the webapp's — it probes the real model on a dummy input to catch a CUDA kernel/compute-capability mismatch that `torch.cuda.is_available()` does not detect, and `classification_curvature`/`load_model` were adjusted to read the device off the model rather than recomputing availability. Webapp's older, unused geometric curvature-profile functions (`compute_curvature_profile`/`compute_curvature`) were intentionally not ported since the webapp itself doesn't wire them into its own pipeline either. |
+| `manual.py` | `manual.py` | `compute_manual_length` | 2026-07-20 | None — module docstring already states this is shared logic between the webapp and this extension. |
+| `scalebar.py` | `scalebar.py` | `detect_scalebar`, `calibrate_from_endpoints`, `draw_scalebar_endpoints` | 2026-07-20 | None known. `calibrate_from_endpoints`/`draw_scalebar_endpoints` exist here but (until issue #76) have no Slicer UI wired up to call them — the webapp exposes them via its manual scale-bar entry accordion. |
+
+## Model presets and weights
+
+Not a core `.py` port, but the same re-sync problem: `ZebrafishEmbryoAnalyzerLib/model_manifest.py`
+mirrors the webapp's `SEG_MODEL_OPTIONS` table. Synced 2026-07-29.
+
+Read that table together with the comment directly above it and with the call
+site in `process()`. **A `None` filename there means "use the pipeline default",
+not "this preset has no such model."** Reading it the other way produced two
+wrong implementations on 2026-07-29 (edema wrongly restricted to the DESY
+preset, and the Fast & Easy body model discarded as a legacy file). Ports of
+this table should be checked against the *resolved* filenames, not the literal
+cell contents.
+
+| Preset (webapp label) | `MODEL_SETS` key | Input size | Notes |
+|---|---|---|---|
+| `Fast & Easy (256 px, ~2s/image)` | `fast` | 256 | Body/eye/edema/swim-bladder all resolve to the pipeline defaults. Its swim bladder model is Unet + vgg16, unlike the 512px presets' FPN + vgg19. |
+| `Complex & Slower (512 px, ~7s/image)` | `general` | 512 | Names body, eye and swim bladder explicitly; edema falls through to the default. |
+| `Fine-tuned DESY` | `desy` | 512 | Names all four explicitly. |
+
+The combo-box labels are the webapp's verbatim so users recognise the same
+presets in both tools. The stable ids (`fast`/`general`/`desy`) are **not** the
+webapp's strings — they are persisted in the MRML parameter node and travel
+inside saved scenes, so they must not be renamed to match a display label.
+
+**Deliberate deviation — edema on the 512px general preset.** The webapp offers
+edema there and, having no 512px general edema model, falls back to the 256px
+default and feeds it 512px input. This port omits the role for that preset
+instead and greys the checkbox out.
+
+Reason: Unet and FPN are fully convolutional and accept any input size divisible
+by 32, so a resolution mismatch raises nothing — it yields a mask from which a
+plausible-looking µm² figure is computed, indistinguishable from a correct one in
+a result table. For a measurement tool a silently wrong number is worse than an
+unavailable feature.
+
+Rollback: if a 512px general edema model appears upstream, add it to `MODELS` and
+wire it into `MODEL_SETS["general"]` under `"edema"`. The checkbox re-enables
+itself — `widget.py` only tests `"edema" in model_set` — and no other change is
+needed.
+
+## Convention for future ports
+
+See `CLAUDE.md` → "Code style" for the naming/signature convention: keep
+ported functions close to the webapp original (same name, same signature)
+unless there's a concrete reason to diverge, and note the reason inline as a
+short comment when there is one. Update the table above whenever a function
+is newly ported or re-synced, including the date and — where practical — a
+pointer to which webapp state (a git SHA, if the webapp source is inspected
+via its git history, or "as deployed on `<date>`" if inspected via the live
+Space) the port was taken from.

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerCore/length.py

@@ -129,6 +129,78 @@ def compute_eye_metrics(mask_eye, mask_fish=None, spacing=(1.0, 1.0)):
         "eye_diameter_points": eye_diameter_points,
     }
 
+def compute_tube_metrics(mask, spacing=(1.0, 1.0)):
+    """
+    Fit a minimum-area rotated rectangle to a tube-shaped binary mask.
+
+    Unlike a simple bounding box, the rectangle follows the tube's actual
+    orientation, so the short side is the cross-sectional width regardless of
+    how the tube is rotated in the image (the long side is discarded — callers
+    that only care about the body length already get that from
+    tube_length_border2border).
+
+    spacing: (dy, dx) physical units per pixel.
+
+    Returns dict with keys:
+        area: physical area (spacing units squared)
+        length: long-axis extent (spacing units) — the "long part"
+        width: short-axis extent (spacing units) — the tube width
+        length_line: ((r1,c1),(r2,c2)) endpoints of the long-axis midline, or None
+        width_line: ((r1,c1),(r2,c2)) endpoints of the width midline, or None
+    """
+    import cv2  # deferred: heavy compiled extension, only needed at call time
+    out = {"area": 0.0, "length": 0.0, "width": 0.0, "length_line": None, "width_line": None}
+    if mask is None:
+        return out
+    m = np.asarray(mask)
+    if m.ndim == 3:
+        m = m[..., 0]
+    m = (m > 0).astype(np.uint8)
+    if not m.any():
+        return out
+    dy, dx = spacing
+
+    num, labels, stats, _ = cv2.connectedComponentsWithStats(m, connectivity=8)
+    if num > 1:
+        largest = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
+        m = (labels == largest).astype(np.uint8)
+
+    out["area"] = float(int(m.sum()) * dy * dx)
+
+    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    if not contours:
+        return out
+    contour = max(contours, key=cv2.contourArea)
+    if len(contour) < 3:
+        return out
+
+    box = cv2.boxPoints(cv2.minAreaRect(contour))  # 4 (x, y) pixel points, in order around the rect
+    box_rc = box[:, ::-1]  # -> (row, col) to match this codebase's point convention
+    A, B, C, D = box_rc
+
+    def _phys_len(p1, p2):
+        dr, dc = (p2[0] - p1[0]) * dy, (p2[1] - p1[1]) * dx
+        return float(np.sqrt(dr ** 2 + dc ** 2))
+
+    def _mid(p1, p2):
+        return ((p1[0] + p2[0]) / 2.0, (p1[1] + p2[1]) / 2.0)
+
+    side_AB, side_BC = _phys_len(A, B), _phys_len(B, C)
+
+    # The segment joining the midpoints of a pair of parallel sides spans the
+    # *other* pair's side length (it cuts straight across the rectangle), so
+    # the line with length == side_AB is mid(BC)-mid(DA), and vice versa.
+    if side_AB >= side_BC:
+        out["length"], out["width"] = side_AB, side_BC
+        out["length_line"] = (_mid(B, C), _mid(D, A))
+        out["width_line"] = (_mid(A, B), _mid(C, D))
+    else:
+        out["length"], out["width"] = side_BC, side_AB
+        out["length_line"] = (_mid(A, B), _mid(C, D))
+        out["width_line"] = (_mid(B, C), _mid(D, A))
+
+    return out
+
 def compute_eye_diameters(mask_eye, spacing=(1.0, 1.0)):
     """
     Measure the horizontal and vertical diameters of the eye from the binary mask.
@@ -658,12 +730,46 @@ def apply_mask(original_image, mask):
 
     return masked_image
 
+def select_torch_device(torch, probe_model=None, probe_input_shape=(1, 3, 256, 256)):
+    """Pick a usable torch device, falling back to CPU when CUDA is unusable.
+
+    torch.cuda.is_available() only confirms a CUDA runtime/driver is present,
+    not that the installed build's compiled kernels cover this GPU's compute
+    capability. On a mismatch, a real kernel launch fails with "CUDA error: no
+    kernel image is available for execution on the device".
+
+    A trivial canary op (e.g. a bare add) is not a reliable stand-in for that
+    check: observed on real hardware to succeed — after a slow one-time CUDA
+    context/JIT warmup — on a GPU where the model's own conv/batchnorm kernels
+    still failed immediately afterwards. Probing with the real model on a dummy
+    input of its expected shape exercises the same kernels the model actually
+    uses, so the failure (if any) shows up here instead of mid-analysis.
+    When probe_model is None (no model to test yet), only is_available() is
+    checked.
+    """
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+    if probe_model is None:
+        return torch.device("cuda")
+    try:
+        probe_model.to("cuda")
+        with torch.no_grad():
+            probe_model(torch.zeros(*probe_input_shape, device="cuda"))
+        return torch.device("cuda")
+    except RuntimeError:
+        probe_model.to("cpu")
+        return torch.device("cpu")
+
+
 def classification_curvature(image, mask, model, use_threshold, threshold):
     import torch
     import torch.nn.functional as F
     import torchvision.transforms as T
     import cv2  # deferred: only needed at call time
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # The model's actual device reflects whatever load_model's CUDA probe decided —
+    # recomputing availability independently here could disagree with it and feed a
+    # cuda tensor to a cpu-fallback model (or vice versa).
+    device = next(model.parameters()).device
 
     masked_image = apply_mask(image, mask)
 
@@ -714,7 +820,6 @@ def load_model(model_path: str):
     import torch.nn as nn
     import torch.nn.functional as F
     import timm
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     import logging as _logging
     _log = _logging.getLogger(__name__)
@@ -752,7 +857,9 @@ def forward(self, x):
     _log.debug("Using curvature model params: %s", best_params)
 
     try:
-        state_dict = torch.load(model_path, map_location=device, weights_only=True)
+        # Always deserialize to CPU first — the final device is decided below, after
+        # the model is built, by actually probing it (see select_torch_device).
+        state_dict = torch.load(model_path, map_location=torch.device("cpu"), weights_only=True)
     except Exception as exc:
         raise RuntimeError(
             f"Failed to load curvature model from {model_path!r} with safe loading. "
@@ -786,9 +893,11 @@ def forward(self, x):
             f"Curvature checkpoint at {model_path!r} is missing required keys: {missing}. "
             "The checkpoint may be incomplete or incompatible. Re-download the model."
         )
-    model = model_instance.to(device)
-    model.eval()
-    return model
+    # Decide the device by actually probing model_instance (see select_torch_device) —
+    # it moves the model to its final device itself, cuda or a cpu fallback.
+    select_torch_device(torch, probe_model=model_instance)
+    model_instance.eval()
+    return model_instance
 
 
 def plot_edges_with_curvature(mask, min_contour_length, window_size_ratio):

ZebrafishEmbryoAnalyzer/ZebrafishEmbryoAnalyzerCore/seg.py

@@ -4,9 +4,9 @@
 
 _UNET_CACHE = {}  # lazy-loaded cache keyed by (filename_or_path, encoder_name)
 
-def _load_unet_model(model_path=None, repo_id=None, filename=None, label="model", revision="main", force_download=False, encoder_name="vgg16"):
+def _load_unet_model(model_path=None, repo_id=None, filename=None, label="model", revision="main", force_download=False, encoder_name="vgg16", model_type="Unet"):
     """
-    Load a binary Unet model from a local path.
+    Load a binary segmentation model from a local path.
 
     ``model_path`` must point to an existing local file.  The ``repo_id``,
     ``filename``, ``revision``, and ``force_download`` parameters are accepted
@@ -15,17 +15,25 @@ def _load_unet_model(model_path=None, repo_id=None, filename=None, label="model"
     ZebrafishEmbryoAnalyzerLib.model_downloader to download model files before
     calling this function.
 
+    ``model_type``: "Unet" (default, used for body/eye/edema) or "FPN" (used for
+    swim bladder — a different segmentation_models_pytorch architecture, not a
+    variant of Unet, so it needs its own constructor call).
+
     Returns the model instance when successful, otherwise None.
     Raises RuntimeError when no usable local path is found.
     """
-    cache_key = (model_path or filename, encoder_name)
+    cache_key = (model_path or filename, encoder_name, model_type)
     if cache_key in _UNET_CACHE:
         print(f"{label.capitalize()} served from cache.")
         return _UNET_CACHE[cache_key]
 
     import torch
-    from segmentation_models_pytorch import Unet
-    model = Unet(encoder_name=encoder_name, encoder_weights="imagenet", in_channels=3, classes=1)
+    if model_type == "FPN":
+        from segmentation_models_pytorch import FPN
+        model = FPN(encoder_name=encoder_name, encoder_weights="imagenet", in_channels=3, classes=1)
+    else:
+        from segmentation_models_pytorch import Unet
+        model = Unet(encoder_name=encoder_name, encoder_weights="imagenet", in_channels=3, classes=1)
     resolved_path = None
 
     if model_path and os.path.exists(model_path):
@@ -79,6 +87,12 @@ def segmentation_pipeline(
     edema_repo_id="markdanielarndt/Zebrafish_Segmentation",
     edema_model_filename="best_model_edema_3400_focal.pth",
     edema_encoder_name="vgg19",
+    include_swimbladder=False,
+    swimbladder_model_path=None,
+    swimbladder_repo_id="markdanielarndt/Zebrafish_Segmentation",
+    swimbladder_model_filename="best_model_swimmbladder_256_09072026.pth",
+    swimbladder_encoder_name="vgg16",
+    swimbladder_model_type="Unet",
 ):
     """
     Perform body segmentation on all images in the specified folder or file list.
@@ -87,12 +101,16 @@ def segmentation_pipeline(
     When `file_list` is provided it takes precedence and preserves the given order.
 
     Optional eye segmentation can be enabled by setting include_eyes=True.
+    Optional swim bladder segmentation can be enabled by setting include_swimbladder=True
+    (uses an FPN model, not Unet — see swimbladder_model_type / _load_unet_model).
     Optional edema segmentation can be enabled by setting include_edema=True.
 
     Returns:
-        - default: (original_images, segmented_images, grown_images)
-        - if include_eyes=True: (original_images, segmented_images, grown_images, eyes_images)
-        - if include_eyes=True and include_edema=True: (original_images, segmented_images, grown_images, eyes_images, edema_images)
+        (original_images, segmented_images, grown_images), followed by
+        eyes_images / edema_images / swimbladder_images — in that fixed order —
+        for whichever of include_eyes / include_edema / include_swimbladder are
+        True. E.g. include_edema=True alone returns a 4-tuple ending in
+        edema_images; all three True returns a 6-tuple.
     """
     import cv2
     if file_list is not None:
@@ -155,6 +173,23 @@ def segmentation_pipeline(
         else:
             print("Edema model loaded successfully!")
 
+    swimbladder_model = None
+    swimbladder_images = []
+    if include_swimbladder:
+        print(f"Loading swim bladder segmentation model from {swimbladder_repo_id}/{swimbladder_model_filename}...")
+        swimbladder_model = _load_unet_model(
+            model_path=swimbladder_model_path,
+            repo_id=swimbladder_repo_id,
+            filename=swimbladder_model_filename,
+            label="swim bladder model",
+            encoder_name=swimbladder_encoder_name,
+            model_type=swimbladder_model_type,
+        )
+        if swimbladder_model is None:
+            print(f"WARNING: Swim bladder model unavailable at {swimbladder_repo_id}/{swimbladder_model_filename}. Returning empty swim bladder masks.")
+        else:
+            print("Swim bladder model loaded successfully!")
+
     import torch
     # Preprocessing parameters
     mean = np.array([0.485, 0.456, 0.406])
@@ -190,14 +225,27 @@ def segmentation_pipeline(
                 segmented_edema_array = np.zeros((target_size[0], target_size[1]), dtype=np.uint8)
             edema_images.append(segmented_edema_array)
 
+        if include_swimbladder:
+            if swimbladder_model is not None:
+                segmented_swimbladder, _ = segment_fish(input_image, swimbladder_model, biggest_only=False)
+                segmented_swimbladder_array = np.array(segmented_swimbladder)
+            else:
+                segmented_swimbladder_array = np.zeros((target_size[0], target_size[1]), dtype=np.uint8)
+            swimbladder_images.append(segmented_swimbladder_array)
+
         grown_images.append(grown_image)
         segmented_images.append(filled_image)
         original_images.append(original_image)
 
-    if include_eyes and include_edema:
-        return original_images, segmented_images, grown_images, eyes_images, edema_images
-
+    # Fixed append order (eyes, edema, swimbladder) regardless of which subset of
+    # the three optional flags is set — the caller (logic.py's analyse_images())
+    # already knows which flags it passed and unpacks accordingly, so an N-way
+    # if/elif tree over 2**3 combinations isn't needed here.
+    result = [original_images, segmented_images, grown_images]
     if include_eyes:
-        return original_images, segmented_images, grown_images, eyes_images
-
-    return original_images, segmented_images, grown_images
+        result.append(eyes_images)
+    if include_edema:
+        result.append(edema_images)
+    if include_swimbladder:
+        result.append(swimbladder_images)
+    return tuple(result)