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