regression

app.py

@@ -77,6 +77,9 @@ HEADS_3D = {
     "kneeMRI",
 }
 
+REGRESSION_HEADS = {
+    "ixi",
+}
 
 DATASET_OPTIONS: Dict[str, str] = {
     "anatomy-ct": "Anatomy CT (test)",
@@ -380,7 +383,7 @@ def load_dataset_sample(
 ) -> Tuple[
     Optional[np.ndarray],
     str,
-    pd.DataFrame,
+    Dict[str, Any],
     Dict[str, Any],
     Optional[Dict[str, Any]],
 ]:
@@ -388,12 +391,11 @@ def load_dataset_sample(
     subset = DEFAULT_DATASET_FOR_HEAD.get(head)
     if not subset:
         gr.Warning("No dataset found for this head.")
-        return None, "", pd.DataFrame(), gr.update(visible=False), None
+        return None, "", gr.update(visible=False), gr.update(visible=False), None
 
     try:
         target_id = parse_target_selection(target_selection)
         image, meta = sample_dataset_example(subset, target_id)
-
         # Apply windowing only for display, keep raw image for model inference
         windowed_image = apply_windowing(image, subset)
         display = to_display_image(windowed_image)
@@ -412,13 +414,13 @@ def load_dataset_sample(
         return (
             display,
             "",  # Reset prediction text
-            pd.DataFrame(),
+            gr.update(visible=False),
             ground_truth_update,
             {"image": image, "mask": meta.get("mask")},  # Store raw image for inference
         )
     except Exception as exc:  # pragma: no cover - surfaced in UI
         gr.Warning(f"Failed to load sample: {exc}")
-        return None, "", pd.DataFrame(), gr.update(visible=False), None
+        return None, "", gr.update(visible=False), gr.update(visible=False), None
 
 
 def format_probabilities(probs: torch.Tensor, id2label: Dict[int, str]) -> pd.DataFrame:
@@ -437,21 +439,25 @@ def format_probabilities(probs: torch.Tensor, id2label: Dict[int, str]) -> pd.Da
 def run_inference(
     image_state: Optional[Dict[str, Any]],
     head: str,
-) -> Tuple[str, pd.DataFrame]:
+) -> Tuple[str, Dict[str, Any]]:
     if not image_state or "image" not in image_state:
-        return "Load a dataset sample or upload an image first.", pd.DataFrame()
+        return "Load a dataset sample or upload an image first.", gr.update(visible=False)
 
     try:
         image = image_state["image"]
-        probs = infer_image(image, head, image_state.get("mask"))
+        output = infer_image(image, head, image_state.get("mask"), return_probs=head not in REGRESSION_HEADS)
+
+        if head in REGRESSION_HEADS:
+            return f"**Prediction:** {output:.3f}", gr.update(visible=False)
 
         # Use id_to_labels.json mapping, fall back to model config if not available
         id2label = load_id_to_labels().get(head, {})
-        df = format_probabilities(probs, id2label)
+
+        df = format_probabilities(output, id2label)
         top_row = df.iloc[0]
         prediction = f"{top_row['label']} (p={top_row['probability']:.3f})"
         result_text = f"**Prediction:** {prediction}"
-        return result_text, df
+        return result_text, gr.update(visible=True, value=df)
     except Exception as exc:  # pragma: no cover - surfaced in UI
         traceback.print_exc()
         return f"Failed to run inference: {exc}", gr.update(visible=False)
@@ -553,7 +559,7 @@ def build_demo() -> gr.Blocks:
                 ground_truth_display = gr.Markdown(visible=False)
                 gr.Markdown("### Predictions")
                 main_prediction = gr.Markdown()
-                prediction_probs = gr.Dataframe(headers=["class_id", "label", "probability"])
+                prediction_probs = gr.Dataframe(headers=["class_id", "label", "probability"], visible=False)
 
         image_state = gr.State()
 

inference.py

@@ -147,6 +147,7 @@ def infer_image(
     image: np.ndarray,
     head: str,
     mask: Any | None = None,
+    return_probs: bool = True,
 ) -> torch.Tensor:
     processor = load_processor()
     model = load_model(head)
@@ -157,5 +158,8 @@ def infer_image(
             processed["mask"] = mask_tensor
         outputs = model(**processed)
         logits = outputs["logits"]
-        probs = torch.nn.functional.softmax(logits[0], dim=-1)
-    return probs
+        if return_probs: 
+            probs = torch.nn.functional.softmax(logits[0], dim=-1)
+            return probs
+        else:
+            return logits[0].squeeze()