Update app.py

app.py

@@ -6,7 +6,6 @@ from models.HybridGNet2IGSC import Hybrid
 from utils.utils import scipy_to_torch_sparse, genMatrixesLungsHeart
 import scipy.sparse as sp
 import torch
-from zipfile import ZipFile
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 hybrid = None
@@ -178,13 +177,6 @@ def removePreprocess(output, info):
     return output
 
 
-def zip_files(files):
-    with ZipFile("complete_results.zip", "w") as zipObj:
-        for idx, file in enumerate(files):
-            zipObj.write(file, arcname=file.split("/")[-1])
-    return "complete_results.zip"
-
-
 def calculate_ctr(landmarks):
     H = landmarks[94:]
     RL = landmarks[0:44]
@@ -192,7 +184,7 @@ def calculate_ctr(landmarks):
     cardiac_width = np.max(H[:, 0]) - np.min(H[:, 0])
     thoracic_width = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
     ctr = cardiac_width / thoracic_width if thoracic_width > 0 else 0
-    return ctr
+    return round(ctr, 3)
 
 
 def segment(input_img):
@@ -220,28 +212,10 @@ def segment(input_img):
     seg_to_save = (outseg.copy() * 255).astype('uint8')
     cv2.imwrite("tmp/overlap_segmentation.png", cv2.cvtColor(seg_to_save, cv2.COLOR_RGB2BGR))
 
-    RL = output[0:44]
-    LL = output[44:94]
-    H = output[94:]
-
-    np.savetxt("tmp/RL_landmarks.txt", RL, delimiter=" ", fmt="%d")
-    np.savetxt("tmp/LL_landmarks.txt", LL, delimiter=" ", fmt="%d")
-    np.savetxt("tmp/H_landmarks.txt", H, delimiter=" ", fmt="%d")
-
-    RL_mask, LL_mask, H_mask = getMasks(output, original_shape[0], original_shape[1])
-
-    cv2.imwrite("tmp/RL_mask.png", RL_mask)
-    cv2.imwrite("tmp/LL_mask.png", LL_mask)
-    cv2.imwrite("tmp/H_mask.png", H_mask)
-
-    zip = zip_files(
-        ["tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt", "tmp/RL_mask.png", "tmp/LL_mask.png",
-         "tmp/H_mask.png", "tmp/overlap_segmentation.png"])
-
     ctr_value = calculate_ctr(output)
+    interpretation = "Normal" if ctr_value < 0.5 else "Cardiomegaly"
 
-    return outseg, ["tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt", "tmp/RL_mask.png",
-                    "tmp/LL_mask.png", "tmp/H_mask.png", "tmp/overlap_segmentation.png", zip], ctr_value
+    return outseg, "tmp/overlap_segmentation.png", ctr_value, interpretation
 
 
 if __name__ == "__main__":
@@ -275,8 +249,12 @@ if __name__ == "__main__":
 
                 with gr.Column():
                     image_output = gr.Image(type="filepath", height=750)
+                    
+                    with gr.Row():
+                        ctr_output = gr.Number(label="CTR (Cardiothoracic Ratio)")
+                        ctr_interpretation = gr.Textbox(label="Interpretation", interactive=False)
+                    
                     results = gr.File()
-                    ctr_output = gr.Number(label="CTR (Cardiothoracic Ratio)")
 
         gr.Markdown("""
                     If you use this code, please cite:
@@ -322,7 +300,8 @@ if __name__ == "__main__":
         clear_button.click(lambda: None, None, image_input, queue=False)
         clear_button.click(lambda: None, None, image_output, queue=False)
         clear_button.click(lambda: None, None, ctr_output, queue=False)
+        clear_button.click(lambda: None, None, ctr_interpretation, queue=False)
 
-        image_button.click(segment, inputs=image_input, outputs=[image_output, results, ctr_output], queue=False)
+        image_button.click(segment, inputs=image_input, outputs=[image_output, results, ctr_output, ctr_interpretation], queue=False)
 
     demo.launch()