import gradio as gr
import torch
from PIL import Image
import json

m_raw_model = torch.hub.load('ultralytics/yolov8', 'custom', path='M-Raw.pt', source="local") 
s_raw_model = torch.hub.load('ultralytics/yolov8', 'custom', path='S-Raw.pt', source="local") 
n_raw_model = torch.hub.load('ultralytics/yolov8', 'custom', path='N-Raw.pt', source="local")
m_pre_model = torch.hub.load('ultralytics/yolov8', 'custom', path='M-Pre.pt', source="local") 
s_pre_model = torch.hub.load('ultralytics/yolov8', 'custom', path='S-Pre.pt', source="local") 
n_pre_model = torch.hub.load('ultralytics/yolov8', 'custom', path='N-Pre.pt', source="local")

def snap(image, model, conf, iou):
    
        # If no model selected, use M-Raw
        if model == None:
            model = "M-Raw"
        
        # Run the selected model
        results = None
        if model == "M-Raw":
            results = m_raw_model(image, conf=conf, iou=iou)
        elif model == "N-Raw":
            results = n_raw_model(image, conf=conf, iou=iou)
        elif model == "S-Raw":
            results = s_raw_model(image, conf=conf, iou=iou)
        elif model == "M-Pre":
            results = m_pre_model(image, conf=conf, iou=iou)
        elif model == "N-Pre":
            results = n_pre_model(image, conf=conf, iou=iou)
        elif model == "S-Pre":
            results = s_pre_model(image, conf=conf, iou=iou)