Update app.py

app.py

@@ -12,6 +12,8 @@ from skimage.filters import threshold_otsu
 import torchio as tio
 
 def infer_full_vol(tensor, model):
+    tensor = tensor.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, D, H, W] - adding batch and channel dims
+    
     tensor = torch.movedim(tensor, -1, -3)
     tensor = tensor / tensor.max()
     
@@ -44,7 +46,7 @@ def infer_full_vol(tensor, model):
     return output.squeeze().detach().cpu().numpy()
 
 def infer_patch_based(tensor, model, patch_size=64, stride_length=32, stride_width=32, stride_depth=16, batch_size=10, num_worker=2):
-    test_subject = tio.Subject(img = tio.ScalarImage(tensor=tensor.unsqueeze(0)))
+    test_subject = tio.Subject(img = tio.ScalarImage(tensor=tensor.unsqueeze(0))) # adding channel dim while creating the TorchIO subject
     overlap = np.subtract(patch_size, (stride_length, stride_width, stride_depth))
     
     with torch.no_grad():
@@ -140,9 +142,6 @@ if uploaded_file is not None and process_button:
             st.write(f"Input tensor shape: `{tensor.shape}`")
             st.write(f"Selected pretrained model: `{selected_model}`")
 
-            # Add batch and channel dimensions
-            tensor = tensor.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, D, H, W]
-
             # Construct the model name based on the selected model
             model_name = f"soumickmj/{selected_model}"