comments

app.py

@@ -616,24 +616,10 @@ def build_demo() -> gr.Blocks:
             outputs=[main_prediction, prediction_probs],
         )
 
-        gr.Markdown(
-            """
-            ### Notes
-
-            - Configure the `HF_TOKEN` secret in your Space to load private checkpoints
-              and datasets from the `raidium` organisation.
-            - When masks are available in the dataset sample, their contours are drawn on the
-              image for visual reference using OpenCV.
-            - Uploaded images must be single-channel arrays. Multi-channel inputs are
-              converted to grayscale automatically.
-            """
-        )
-
     return demo
 
 
 demo = build_demo()
 
-
 if __name__ == "__main__":
     demo.launch()

inference.py

@@ -82,11 +82,9 @@ def prepare_mask_for_model(mask: Any) -> Optional[torch.Tensor]:
     if mask_arr.size == 0:
         return None
 
-    if mask_arr.ndim == 3:
-        tensor = mask_transform(mask_arr.transpose(2, 0, 1))
-        # Match the shape produced in simple_test.py so the model receives
-        # (batch, height, width, channels) style tensors.
-        tensor = tensor.transpose(1, 3).transpose(1, 2)
+    if mask_arr.ndim == 3: # (H, W, slices)
+        tensor = mask_transform(mask_arr.transpose(2, 0, 1)) # (1, slices, H, W)
+        tensor = tensor.transpose(1, 3).transpose(1, 2) # 
     else:
         tensor = mask_transform(torch.tensor([mask_arr]))
         tensor = tensor.unsqueeze(0)