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")