import gradio as gr
from ultralytics import YOLO
import cv2

# Define the paths to the weights files
MODEL_PATHS = {
    "YOLO - 100 Epochs": "100best.pt",
    "YOLO - 150 Epochs": "150best.pt",
    "YOLO - 200 Epochs": "200best.pt"
}

def detect_and_count(image, model_choice):
    # Dynamically load the weights for the selected choice
    weights_path = MODEL_PATHS[model_choice]
    selected_model = YOLO(weights_path)
    
    # Run prediction
    results = selected_model.predict(source=image, conf=0.30, iou=0.45)
    result = results[0]
    
    # Extract structural analytics
    object_count = len(result.boxes)
    annotated_image = result.plot()
    
    summary_text = f"Active Weights: {weights_path}\nTotal objects detected: {object_count}"
    return annotated_image, summary_text

# Define the user interface
interface = gr.Interface(
    fn=detect_and_count,
    inputs=[
        gr.Image(type="numpy", label="Upload an Image"),
        gr.Dropdown(
            choices=["YOLO - 100 Epochs", "YOLO - 150 Epochs", "YOLO - 200 Epochs"], 
            value="YOLO - 100 Epochs", 
            label="Select Training Duration (Epochs)"
        )
    ],
    outputs=[
        gr.Image(type="numpy", label="Detection Output"),
        gr.Textbox(label="Performance Summary")
    ],
    title="Multi-Epoch Comparison Object Detector",
    description="Upload an image and switch between the 100, 150, and 200 epoch models. Click 'Submit' after switching to update."
)

if __name__ == "__main__":
    interface.launch()