Added model and threshold choosing

app.py

@@ -1,30 +1,57 @@
+import os
 import gradio as gr
 from ultralytics import YOLO
 import numpy as np
-import os
-
-# Load YOLO model
-model = YOLO('./best.pt')
 
+model_options = ["yolo-8n-dota.pt", "yolo-8s-dota.pt", "yolo-8m-dota.pt"]
+model_names = ["Nano", "Small", "Medium"]
+models = [YOLO(option) for option in model_options]
 example_list = [["examples/" + example] for example in os.listdir("examples")]
 
-def process_image(input_image):
-    if input_image is not None:
-        results = model(input_image)
+def process_image(input_image, model_name, conf):
+    if input_image is None:
+        return None, "No image"
+    
+    if model_name is None:
+        model_name = model_names[0]
+    
+    if conf is None:
+        conf = 0.6
+
+    model_index = model_names.index(model_name)
+    model = models[model_index]
+    
+    results = model.predict(input_image, conf=conf)
+    class_counts = {}
+    class_counts_str = "Class Counts:\n"
+
+    for r in results:
+        im_array = r.plot()
+        im_array = im_array.astype(np.uint8)
+
+        if len(r.obb.cls) == 0:  # If no objects are detected
+            return None, "No objects detected."
+
+        for cls in r.obb.cls:
+            class_name = r.names[cls.item()]
+            class_counts[class_name] = class_counts.get(class_name, 0) + 1
+
+        for cls, count in class_counts.items():
+            class_counts_str += f"\n{cls}: {count}"
 
-        for r in results:
-            im_array = r.plot()
-            im_array = im_array.astype(np.uint8)
-            return im_array
+        return im_array, class_counts_str
 
-# Create Gradio Interface
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(),
-    outputs=gr.Image(),  # Specify output as Gradio Image
+    inputs=[
+        gr.Image(), 
+        gr.Radio(model_names, label="Choose model", value=model_names[0]),
+        gr.Slider(minimum=0.2, maximum=1.0, step=0.1, label="Confidence Threshold", value=0.6)
+    ], 
+    outputs=["image", gr.Textbox(label="More info")],
     title="YOLOv8-obb aerial detection",
-    description="YOLOv8-obb trained on DOTAv1.5",
-    examples=example_list)
+    description='''YOLOv8-obb trained on DOTAv1.5''',
+    examples=example_list
+)
 
-# Launch the Gradio interface
-iface.launch()
+iface.launch()

flagged/log.csv

@@ -1,2 +0,0 @@
-name,intensity,output,flag,username,timestamp
-aasd,99,"Hello, aasd!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!",,,2024-02-20 13:01:37.633520

yolo-8n-dota.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ee44108137e10e13a377e2a75175af1476c355e04217cc38e3e8e2f4cb6fd7c
+size 6465538

yolo-8s-dota.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4395070f75af8464e3c4d7e7d83eda61d19d6cd5ea6c62ca1898844a8e0ad54c
+size 23169282