You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
94 lines
3.7 KiB
94 lines
3.7 KiB
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[0] # YOLOv5 root directory
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
|
|
|
from models.common import DetectMultiBackend
|
|
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
|
|
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
|
|
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
|
|
from utils.torch_utils import select_device, time_sync
|
|
|
|
|
|
class Detector:
|
|
@torch.no_grad()
|
|
def __init__(self):
|
|
# Load model
|
|
weights = ROOT / 'best.pt' # model.pt path(s)
|
|
data = ROOT / 'maskhelper.yaml' # dataset.yaml path
|
|
imgsz = (640, 640) # inference size (height, width)
|
|
self.device = select_device('')
|
|
self.model = DetectMultiBackend(weights, device=self.device, dnn=False, data=data, fp16=False)
|
|
self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
|
|
self.imgsz = check_img_size(imgsz, s=self.stride) # check image size
|
|
|
|
|
|
def detect(self,
|
|
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam,
|
|
conf_thres=0.85, # confidence threshold
|
|
iou_thres=0.85, # NMS IOU threshold
|
|
max_det=10, # maximum detections per image
|
|
classes=None, # filter by class: --class 0, or --class 0 2 3
|
|
agnostic_nms=False, # class-agnostic NMS
|
|
augment=False, # augmented inference
|
|
):
|
|
self.source = str(source)
|
|
# Dataloader
|
|
dataset = LoadImages(self.source, img_size=self.imgsz, stride=self.stride, auto=self.pt)
|
|
self.bs = 1 # batch_size
|
|
# Run inference
|
|
self.model.warmup(imgsz=(1 if self.pt else self.bs, 3, *self.imgsz)) # warmup
|
|
seen, windows, dt = 0, [], [0.0, 0.0, 0.0]
|
|
for path, im, im0s, vid_cap, s in dataset:
|
|
t1 = time_sync()
|
|
im = torch.from_numpy(im).to(self.device)
|
|
im = im.half() if self.model.fp16 else im.float() # uint8 to fp16/32
|
|
im /= 255 # 0 - 255 to 0.0 - 1.0
|
|
if len(im.shape) == 3:
|
|
im = im[None] # expand for batch dim
|
|
t2 = time_sync()
|
|
dt[0] += t2 - t1
|
|
|
|
# Inference
|
|
pred = self.model(im, augment=augment, visualize=False)
|
|
t3 = time_sync()
|
|
dt[1] += t3 - t2
|
|
|
|
# NMS
|
|
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
|
|
dt[2] += time_sync() - t3
|
|
|
|
# Second-stage classifier (optional)
|
|
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
|
|
|
|
# Process predictions
|
|
for i, det in enumerate(pred): # per image
|
|
res = []
|
|
seen += 1
|
|
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
|
|
|
|
s += '%gx%g ' % im.shape[2:] # print string
|
|
|
|
if len(det):
|
|
# Rescale boxes from img_size to im0 size
|
|
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
|
|
|
|
# Write results
|
|
for *xyxy, conf, cls in reversed(det):
|
|
c = int(cls) # integer class
|
|
label = self.names[c]
|
|
coords = [xyxy[0], xyxy[1], xyxy[2], xyxy[3]]
|
|
|
|
res.append([label,coords])
|
|
return res
|
|
|