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.detach().cpu().numpy()
+    return output.squeeze().detach().cpu().numpy()
 
 # Set page configuration
 st.set_page_config(
@@ -35,7 +35,7 @@ st.set_page_config(
 with st.sidebar:
     st.title("Segmenting vessels in the brain from a 3D Magnetic Resonance Angiograph, ideally acquired at 7T | DS6")
     st.markdown("""
-    This application allows you to upload a 3D NIfTI file (dims: H x W x D), process it through a pre-trained 3D model (from DS6 and other related works), and download the output as a `.nii.gz` file containing the vessel segmentation.
+    This application allows you to upload a 3D NIfTI file (dims: H x W x D, where the final dim is the slice dim in the axial plane), process it through a pre-trained 3D model (from DS6 and other related works), and download the output as a `.nii.gz` file containing the vessel segmentation.
     
     **Instructions**:
     - Upload your 3D NIfTI file (`.nii` or `.nii.gz`). It should be a single-slice cardiac long-axis dynamic CINE scan, where the first dimension represents time.
@@ -46,7 +46,7 @@ with st.sidebar:
     st.markdown("© 2024 Soumick Chatterjee")
 
 # Main content
-st.header("From single-slice cardiac long-axis dynamic CINE scan (3D: HxWxD) to 128 latent factors...")
+st.header("DS6, Deformation-Aware Semi-Supervised Learning: Application to Small Vessel Segmentation with Noisy Training Data")
 
 # File uploader
 uploaded_file = st.file_uploader(