Update app.py

app.py

@@ -1,10 +1,5 @@
 ##importing the libraries
-
-import torch
-import cv2
-import os
-import numpy as np
-import pandas as pd
+from PIL import Image
 import gradio as gr
 from ultralytics import YOLO
 
@@ -12,36 +7,20 @@ from ultralytics import YOLO
 # Load your trained model
 model = YOLO('best.pt')
 
-### Preprocess the new image
-
-def predict_image(test_image):
-    # save images
-    model.predict(test_image, save=True, conf=0.1, batch=2, show=True)
-
-    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]  # getting the latest predictions
-    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
+#Function for making predictions
+def predict (image):
+    results = model(image)
+    for result in results:
+        im_array = result.plot()
+        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
+        # im.show()  # show image
+        # im.save('results.jpg')
+    return im
         
 
 
 
-platform = gr.Interface( fn = predict_image, 
+platform = gr.Interface( fn = predict, 
                         title ="PTCADx: Computer-Aided Detection of Pneumothorax in Chest X-ray Images",
                         inputs = "image", 
                         outputs = "image",