Spaces:
Running
Running
File size: 2,201 Bytes
58e7ec3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import numpy as np
import torch
from models.yolo.utils.datasets import letterbox
from models.yolo.utils.general import check_img_size, non_max_suppression, scale_coords
from models.yolo.utils.torch_utils import select_device
from models.yolo.model import Model
from models.yolo.experimental import attempt_load
class Detector:
def __init__(
self,
model_file,
conf_thres = 0.25,
iou_thres = 0.45,
device = "cpu"
):
self.device = select_device(device)
weights = torch.load(model_file, self.device)
self.model = weights["model"]
self.stride = int(self.model.stride.max())
self.imgsz = check_img_size(640, s=self.stride)
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.half = device != "cpu"
if self.half:
self.model.half()
self.model.eval()
def __call__(
self,
img
):
shape = img.shape
img = letterbox(img, self.imgsz, stride=self.stride)[0] # resize & padding
img = img[:, :, ::-1].transpose(2, 0, 1) # hwc(bgr) -> hwc(rgb) -> c(rgb)hw
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() if self.half else img.float()
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
with torch.no_grad():
pred = self.model(img, augment=True)[0]
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)[0]
if not len(pred):
return
# Scale boxes size back
pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], shape).round()
pred_ = pred[:, [0, 1, 2, 3, 5]].to("cpu", int).numpy()
cls = np.unique(pred_[:, -1], return_index=False)
if 0 in cls:
code = "ssq"
elif 2 in cls:
code = "cjdlt"
else:
return
if not 1 in cls and 3 in cls:
return
numbers = pred_[pred_[:, -1] == 1]
issue = pred_[pred_[:, -1] == 3]
return code, issue, numbers
|