Update main.py

main.py

@@ -86,17 +86,17 @@ class ProstateSegmentationAlgorithm(SegmentationAlgorithm):
 
         # input / output paths for algorithm
         self.input_dirs = [
-            "/input/images/transverse-t2-prostate-mri"
+            "./input/images/transverse-t2-prostate-mri"
         ]
         self.scan_paths = []
-        self.prostate_segmentation_path_pz = Path("/output/images/softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz.mha")
-        self.prostate_segmentation_path_tz = Path("/output/images/softmax-prostate-central-gland-segmentation/prostate_gland_sm_tz.mha")
-        self.prostate_segmentation_path = Path("/output/images/prostate-zonal-segmentation/prostate_gland.mha")
+        self.prostate_segmentation_path_pz = Path("./output/images/softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz.mha")
+        self.prostate_segmentation_path_tz = Path("./output/images/softmax-prostate-central-gland-segmentation/prostate_gland_sm_tz.mha")
+        self.prostate_segmentation_path = Path("./output/images/prostate-zonal-segmentation/prostate_gland.mha")
 
         # input / output paths for nnUNet
-        self.nnunet_inp_dir = Path("/opt/algorithm/nnunet/input")
-        self.nnunet_out_dir = Path("/opt/algorithm/nnunet/output")
-        self.nnunet_results = Path("/opt/algorithm/results")
+        self.nnunet_inp_dir = Path("./nnunet/input")
+        self.nnunet_out_dir = Path("./nnunet/output")
+        self.nnunet_results = Path("./results")
 
         # ensure required folders exist
         self.nnunet_inp_dir.mkdir(exist_ok=True, parents=True)
@@ -246,16 +246,24 @@ class ProstateSegmentationAlgorithm(SegmentationAlgorithm):
 def predict(input_file):
     print("Making prediction")
     image = sitk.ReadImage(input_file)
-    print(image.size)
+    sitk.WriteImage(image, "./input/images/transverse-t2-prostate-mri/1009_2222_t2w.mha")
+    
     ProstateSegmentationAlgorithm().process()
-    return input_file
+
+    
+    return (
+        "./output/images/softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz.mha",
+        "./output/images/softmax-prostate-central-gland-segmentation/prostate_gland_sm_tz.mha",
+        "./output/images/prostate-zonal-segmentation/prostate_gland.mha"
 
 print("Starting interface")
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.File(file_count="single", file_types=[".mha", ".nii.gz", ".nii"]),
+    inputs=gr.File(label="input T2 image (3d)", file_count="single", file_types=[".mha", ".nii.gz", ".nii"]),
     outputs=(
-        gr.File()
+        gr.File(label="softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz"),
+        gr.File(label="softmax-prostate-central-gland-segmentation/prostate_gland_sm_tz"),
+        gr.File(label="prostate-zonal-segmentation/prostate_gland"),
     ),
     cache_examples=False,
     # outputs=gr.Label(num_top_classes=3),