Fix /mask-data crash for fresh auto-segment sessions

flask-server/api/api_blueprint.py

@@ -645,18 +645,27 @@
     if not session_key:
         return jsonify({"error": "Missing sessionKey"}), 400
 
-    result = get_mask_data_internal(session_key)
-
-    # For numeric PanTS cases, fall back to the robust label-based computation
-    # (with HuggingFace download) when the local dataset path is unavailable,
-    # so organ statistics resolve even without a full local PanTS install.
+    # get_mask_data_internal() is PanTS-catalog-only (int(id) is its first
+    # line): calling it with a fresh-upload session's UUID always raised and
+    # was silently swallowed into an {"error": ...} dict, so the viewer never
+    # got real organ stats for an auto-segmented scan. Route those through the
+    # session-based lookup instead; only numeric catalog ids still go through
+    # get_mask_data_internal.
     if str(session_key).strip().isdigit():
+        result = get_mask_data_internal(session_key)
+
+        # For numeric PanTS cases, fall back to the robust label-based computation
+        # (with HuggingFace download) when the local dataset path is unavailable,
+        # so organ statistics resolve even without a full local PanTS install.
         if not isinstance(result, dict) or result.get("error") or not result.get("organ_metrics"):
             robust = _ai_compute_organ_metrics_from_labels(str(session_key).strip())
             if robust and robust.get("organ_metrics"):
                 result = robust
+    else:
+        job = _get_inference_job(session_key)
+        result = get_session_mask_data(session_key, job)
 
     return jsonify(result)
 
 
 @api_blueprint.route('/get-main-nifti/<clabel_id>.nii.gz', methods=['GET'])

flask-server/api/utils.py

@@ -178,6 +178,88 @@ def get_mask_data_internal(id, fallback=False):
         print(f"[ERROR] get_mask_data_internal: {e}")
         return {"error": str(e)}
 
+
+def get_session_mask_data(session_id, job):
+    """Organ metadata for a freshly-uploaded/auto-segmented scan, keyed by the
+    session's own UUID rather than a PanTS catalog id.
+
+    get_mask_data_internal() assumes a numeric PanTS catalog id (int(id) is the
+    first thing it does) and has no fallback for a session id, so every fresh
+    auto_segment result hit that exception and the viewer never saw real organ
+    stats -- only whatever the frontend showed in the absence of any data.
+    This reuses the session's own completed-job record (ct_path,
+    output_mask_dir, model) instead of the PanTS catalog, and picks the label
+    map for whichever model actually produced the segmentation.
+
+    `job` is the dict api_blueprint._get_inference_job(session_id) returns --
+    the in-memory/disk-mirrored auto_segment job record (SESSIONS_DIR =
+    .../tmp), not services.job_store's DB-backed job (that table has no rows
+    for this code path -- verified directly against the real DB, not assumed).
+    Caller passes it in instead of this module doing its own lookup, since
+    _get_inference_job lives in api_blueprint.py and importing it back here
+    would be circular (api_blueprint already does `from .utils import *`).
+    """
+    from services.auto_segmentor import (
+        _EPAI_TO_VIEWER, _ATLASNET_TO_VIEWER, _SUPREM_TO_VIEWER,
+        _MEDIA_AGENTIC_ORGANS_TO_VIEWER, _MEDIA_AGENTIC_VERTEBRAE_TO_VIEWER,
+        _LESIONSEG_TO_VIEWER, _VIEWER_LABELS,
+    )
+    # _VIEWER_LABELS is {name: viewer_int}; the *_TO_VIEWER maps are
+    # {model_raw_label_int: viewer_int}. calculate_metrics() needs
+    # {name: model_raw_label_int} -- it indexes the combined-labels volume
+    # (which stores the model's own raw ints) directly by that value, so the
+    # viewer_int is only useful here as the join key back to a readable name.
+    _viewer_int_to_name = {v: k for k, v in _VIEWER_LABELS.items()}
+
+    model_to_viewer_map = {
+        "ePAI": _EPAI_TO_VIEWER,
+        "Atlas-Net": _ATLASNET_TO_VIEWER,
+        "SuPreM": _SUPREM_TO_VIEWER,
+        "MedIA-Agentic-Organs": _MEDIA_AGENTIC_ORGANS_TO_VIEWER,
+        "MedIA-Agentic-Vertebrae": _MEDIA_AGENTIC_VERTEBRAE_TO_VIEWER,
+        "LesionSegmenter": _LESIONSEG_TO_VIEWER,
+    }
+
+    try:
+        if not job:
+            return {"error": f"No job found for session {session_id}"}
+        if job.get("status") != "completed":
+            return {"error": f"Job for session {session_id} is not completed (status={job.get('status')!r})"}
+
+        ct_path = job.get("ct_path")
+        output_mask_dir = job.get("output_mask_dir")
+        model = job.get("model")
+        if not ct_path or not output_mask_dir or not model:
+            return {"error": "Completed job is missing ct_path, output_mask_dir, or model"}
+
+        viewer_map = model_to_viewer_map.get(model)
+        if viewer_map is None:
+            return {"error": f"No organ label map known for model {model!r}"}
+
+        combined_labels_path = os.path.join(output_mask_dir, Constants.COMBINED_LABELS_NIFTI_FILENAME)
+        if not os.path.exists(ct_path) or not os.path.exists(combined_labels_path):
+            return {"error": "Session output files are missing on disk"}
+
+        # Multiple raw labels can share one viewer name (e.g. ePAI's three
+        # lesion subtypes all report as "pancreatic_lesion"); later entries
+        # win, same as every other place in this codebase that inverts one of
+        # these maps. Good enough to show real per-organ stats instead of
+        # nothing; not a fix for that pre-existing many-to-one collapse.
+        organ_intensities = {}
+        for model_label_val, viewer_int in viewer_map.items():
+            name = _viewer_int_to_name.get(viewer_int)
+            if name:
+                organ_intensities[name] = model_label_val
+
+        nifti_processor = NiftiProcessor(ct_path, combined_labels_path, organ_intensities)
+        organ_metadata = nifti_processor.calculate_metrics()
+        return clean_nan(organ_metadata)
+
+    except Exception as e:
+        print(f"[ERROR] get_session_mask_data: {e}")
+        return {"error": str(e)}
+
+
 def generate_distinct_colors(n):
     """Generate n visually distinct RGB colors."""
     import colorsys