Added the YOLO Model

app.py

@@ -9,7 +9,7 @@ from bulk_bulge_generation import definitions, smooth
 import fastai
 from fastcore.all import *
 from fastai.vision.all import *
-from ultralytics import YOLO
+from ultralytics import ASSETS, YOLO
 import cv2
 
 def apply_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20):
@@ -277,6 +277,23 @@ learn_fresh = load_learner('model_fresh.pkl')
 # Loads the YOLO Model
 model = YOLO("bulge_yolo_model.pt")
 
+def predict_image(img, conf_threshold, iou_threshold):
+    """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds."""
+    results = model.predict(
+        source=img,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        show_labels=True,
+        show_conf=True,
+        imgsz=640,
+    )
+
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
+
+    return im
+
 def transform_image(image, func_choice, randomization_check, radius, center_x, center_y, strength, reverse_gradient=True, spiral_frequency=1):
     I = np.asarray(Image.open(image))    
 
@@ -414,7 +431,9 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
 
     result_localization = model.predict(transformed, save=True)
 
-    return transformed, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
+    YOLO_image = predict_image(transformed, 0.5, 0.5)
+
+    return transformed, YOLO_image, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
 
 
 demo = gr.Interface(
@@ -439,6 +458,7 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(label="Transformed Image"),
+        gr.Image(label="YOLO Classification"),
         gr.Label(),
         gr.Label(),
         gr.Image(label="Gradient Vector Field"),