|
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 model == None: |
|
model = "M-Raw" |
|
|
|
|
|
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) |
|
|
|
|