Update app.py

app.py

@@ -1,25 +1,46 @@
 import gradio as gr
 import PIL.Image as Image
 from ultralytics import YOLO
+import torch
+import os
 
-model = None
+# Load model once at startup
+print("Loading YOLOv8m Defence model...")
+model = YOLO("yolov8m_defence.pt")
 
-def predict_image(img, conf_threshold, iou_threshold, model_name):
+# Set device and optimize for CPU inference
+if torch.cuda.is_available():
+    device = 'cuda'
+    print("Using GPU acceleration")
+else:
+    device = 'cpu'
+    print("Using CPU inference")
+
+model.to(device)
+
+def predict_image(img, conf_threshold, iou_threshold):
     """Predicts objects in an image using YOLOv8m Defence model with adjustable confidence and IOU thresholds."""
-    global model
-    model = YOLO(model_name)
-    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
+    try:
+        results = model.predict(
+            source=img,
+            conf=conf_threshold,
+            iou=iou_threshold,
+            show_labels=True,
+            show_conf=True,
+            imgsz=640,
+            verbose=False,
+            device=device
+        )
+        
+        for r in results:
+            im_array = r.plot()
+            im = Image.fromarray(im_array[..., ::-1])
+        
+        return im
+    
+    except Exception as e:
+        print(f"Error during prediction: {e}")
+        return img
 
 # Custom CSS for font styling
 css = """
@@ -39,13 +60,13 @@ body, .gradio-container {
 }
 """
 
+# Create interface
 iface = gr.Interface(
     fn=predict_image,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
-        gr.Radio(choices=["yolov8m_defence.pt"], label="Model Name", value="yolov8m_defence.pt"),
     ],
     outputs=gr.Image(type="pil", label="Detection Results"),
     title="YOLOv8m Defence Object Detection",
@@ -56,21 +77,25 @@ iface = gr.Interface(
     Vehicles (car, truck, tank, bus, van), Ships (cargo, yacht, cruise, warship, sailboat), 
     and specialized items (drone, missile).
     
+    **Note:** Running on the Free Tier - inference may take longer than expected.
+    
     Developed for DSTA Brainhack 2025 - TIL-AI Category (Semi-Finalist)
     """,
     examples=[
-        ["examples/test1.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test2.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test3.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test4.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test5.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test6.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test7.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test8.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test9.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
-        ["examples/test10.jpg", 0.25, 0.45, "yolov8m_defence.pt"],
+        ["examples/test1.jpg", 0.25, 0.45],
+        ["examples/test2.jpg", 0.25, 0.45],
+        ["examples/test3.jpg", 0.25, 0.45],
+        ["examples/test4.jpg", 0.25, 0.45],
+        ["examples/test5.jpg", 0.25, 0.45],
+        ["examples/test6.jpg", 0.25, 0.45],
+        ["examples/test7.jpg", 0.25, 0.45],
+        ["examples/test8.jpg", 0.25, 0.45],
+        ["examples/test9.jpg", 0.25, 0.45],
+        ["examples/test10.jpg", 0.25, 0.45],
     ],
-    css=css
+    css=css,
+    cache_examples=True
 )
 
-iface.launch(share=True)
+if __name__ == "__main__":
+    iface.launch(share=True)