Update main.py

main.py

@@ -252,19 +252,24 @@ class csPCaAlgorithm(SegmentationAlgorithm):
         print(subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=True).stdout)
 
 
-def predict(input_file):
+def predict(t2_file, adc_file, hbv_file):
     print("Making prediction")
-    image = sitk.ReadImage(input_file)
+    t2_file = sitk.ReadImage(t2_file)
+    adc_file = sitk.ReadImage(adc_file)
+    hbv_file = sitk.ReadImage(hbv_file)
     
     os.makedirs("./input/images/transverse-t2-prostate-mri/", exist_ok=True)
+    os.makedirs("./input/images/transverse-adc-prostate-mri/", exist_ok=True)
+    os.makedirs("./input/images/transverse-hbv-prostate-mri/", exist_ok=True)
     os.makedirs("./output/images/softmax-prostate-peripheral-zone-segmentation", exist_ok=True)
     os.makedirs("./output/images/softmax-prostate-central-gland-segmentation", exist_ok=True)
     os.makedirs("./output/images/prostate-zonal-segmentation", exist_ok=True)
     
-    sitk.WriteImage(image, "./input/images/transverse-t2-prostate-mri/1009_2222_t2w.mha")
+    sitk.WriteImage(t2_file, "./input/images/transverse-t2-prostate-mri/1009_2222_t2w.mha")
+    sitk.WriteImage(adc_file, "./input/images/transverse-t2-prostate-mri/1009_2222_adc.mha")
+    sitk.WriteImage(hbv_file, "./input/images/transverse-t2-prostate-mri/1009_2222_hbv.mha")
     
     csPCaAlgorithm().process()
-
     
     return (
         "./output/images/softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz.mha",
@@ -278,7 +283,12 @@ demo = gr.Interface(
     description="description text",
     article="article text",
     fn=predict,
-    inputs=gr.File(label="input T2 image (3d)", 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"]),
+        gr.File(label="input ADC image (3d)", file_count="single", file_types=[".mha", ".nii.gz", ".nii"]),
+        gr.File(label="input HBV image (3d)", file_count="single", file_types=[".mha", ".nii.gz", ".nii"]),
+    ],
+        
     outputs=[
         gr.File(label="softmax-prostate-peripheral-zone-segmentation/prostate_gland_sm_pz"),
         gr.File(label="softmax-prostate-central-gland-segmentation/prostate_gland_sm_tz"),