Create model.py

model.py

@@ -0,0 +1,32 @@
+import os
+import gdown
+from ultralytics import YOLO
+import torch
+import streamlit as st
+
+def download_pretrained_model():
+    """Download and store pre-trained model if not present."""
+    model_dir = 'pretrained_models'
+    model_path = os.path.join(model_dir, 'yolov8n.pt')
+    os.makedirs(model_dir, exist_ok=True)
+
+    if not os.path.exists(model_path):
+        st.info("Downloading pre-trained model...")
+        url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt'
+        gdown.download(url, model_path, quiet=False)
+    return model_path
+
+class AircraftDetector:
+    def __init__(self, use_pretrained=True):
+        """Initialize the detector with selected model."""
+        custom_model_path = 'runs/detect/aircraft_detection/weights/best.pt'
+        pretrained_model_path = download_pretrained_model()
+
+        if use_pretrained or not os.path.exists(custom_model_path):
+            st.warning("Using pre-trained YOLOv8 model.")
+            self.model = YOLO(pretrained_model_path)
+        else:
+            st.success("Using custom-trained YOLOv8 model.")
+            self.model = YOLO(custom_model_path)
+        
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")