Update app.py

app.py

@@ -16,8 +16,26 @@ model = YOLO('best.pt')
 
 def predict_image(test_image):
     # save images
-    results = model.predict(test_image, save=True, conf=0.1, batch=2)
-    return results
+    model.predict(test_image, save=True, conf=0.1, batch=2)
+
+    path = "runs\\detect"
+    detection = (os.listdir(path))
+
+    def extract_digits(x):
+        digits = ''.join(filter(str.isdigit, x))
+        return int(digits) if digits else float('inf')  #Use float('inf') as a default value
+
+    sorted_ = sorted(detection, key=extract_digits)
+
+    if len(sorted_)>1:
+        latest_detection = sorted_[-2]  #
+    else:
+        latest_detection = sorted_[0]
+
+    detection_path = os.path.join(path, latest_detection)
+
+    pred = cv2.imread(os.path.join(detection_path, test_image))
+    return pred
         
 
 
@@ -25,7 +43,7 @@ def predict_image(test_image):
 platform = gr.Interface( fn = predict_image, 
                         title ="PTCADx: Computer-Aided Detection of Pneumothorax in Chest X-ray Images",
                         inputs = "image", 
-                        outputs = "annotedimage",
+                        outputs = "image",
                         description="""
                         Introducing a revolutionary computer-aided detection tool designed to enhance the efficiency of clinicians in detecting pneumothorax in chest X-ray images. 
                         """,