Update app.py

app.py

@@ -21,7 +21,7 @@ def infer_full_vol(tensor, model):
             output = output[0]
         output = torch.sigmoid(output)
     output = torch.movedim(output, -3, -1).type(tensor.type())
-    return output.squeeze(0)
+    return output.detach().cpu().numpy()
 
 # Set page configuration
 st.set_page_config(
@@ -134,12 +134,9 @@ if uploaded_file is not None and process_button:
                 print(error)
                 output = output > 0.5  # exception only if input image seems to have just one color 1.0.
             output = output.astype('uint16')
-
-            # Convert output to NumPy array
-            output_np = output.detach().cpu().numpy()
             
             # Save the output as a NIfTI file
-            output_img = nib.Nifti1Image(output_np, affine=nifti_img.affine)
+            output_img = nib.Nifti1Image(output, affine=nifti_img.affine)
             output_path = tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False).name
             nib.save(output_img, output_path)