Update app.py

app.py

@@ -275,9 +275,11 @@ learn_bias = load_learner('model_bias.pkl')
 learn_fresh = load_learner('model_fresh.pkl')
 
 # Loads the YOLO Model
-model = YOLO("bulge_yolo_model.pt")
+model_bulge = YOLO("bulge_yolo_model.pt")
+modelv8x = YOLO("yolov8x.pt")
+modelv8n = YOLO("yolov8n.pt")
 
-def predict_image(img, conf_threshold, iou_threshold):
+def predict_image(img, model, conf_threshold, iou_threshold):
     """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds."""
     results = model.predict(
         source=img,
@@ -431,9 +433,11 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
 
     result_localization = model.predict(transformed, save=True)
 
-    YOLO_image = predict_image(transformed, 0.5, 0.5)
+    YOLO_image = predict_image(transformed, model_bulge, 0.5, 0.5)
+    YOLO_image1 = predict_image(transformed, modelv8n, 0.5, 0.5)
+    YOLO_image2 = predict_image(transformed, modelv8x, 0.5, 0.5)
 
-    return transformed, YOLO_image, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
+    return transformed, YOLO_image, YOLO_image1, YOLO_image2, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
 
 
 demo = gr.Interface(
@@ -458,7 +462,9 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(label="Transformed Image"),
-        gr.Image(label="YOLO Classification"),
+        gr.Image(label="bulge_model Model Classification"),
+        gr.Image(label="yolov8n Model Classification"),
+        gr.Image(label="yolov8x Model Classification"),
         gr.Label(),
         gr.Label(),
         gr.Image(label="Gradient Vector Field"),