import time
import numpy as np
import torch
from models.common import DetectMultiBackend
from utils.augmentations import letterbox
from utils.general import (check_img_size, cv2,
non_max_suppression, scale_coords)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device
class YOLOv5Detector:
""" YOLOv5 object detection """
def __init__(self, weights='yolov5s.pt', conf_thres=0.25, iou_thres=0.45, imgsz=640, data='data/coco128.yaml'):
""" Initialization """
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.device = select_device('0')
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 image_preprocess(self, image):
im0 = image.copy()
im = letterbox(im0, self.imgsz, stride=32, auto=True)[0] # padded resize
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
im = np.ascontiguousarray(im) # contiguous
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
# Dataloader
return im
def __call__(self, image, *args, **kwargs):
im = self.image_preprocess(image)
pred = self.model(im, augment=False, visualize=False)
pred = non_max_suppression(pred,
conf_thres=0.25,
iou_thres=0.45,
classes=None,
agnostic=False,
multi_label=False,
labels=(),
max_det=1000)
for i, det in enumerate(pred): # per image
annotator = Annotator(im0, example=str(self.names))
if len(det):
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
for *xyxy, conf, cls in reversed(det):
label = f'{self.names[int(cls)]} {conf:.2f}'
annotator.box_label(xyxy, label, color=colors(2, True))
return im0
yolov5_detector = YOLOv5Detector(weights='best.pt')
img = r'C:\Users\Administrator\Desktop\000000011244.jpg'
while True:
im0 = cv2.imread(img)
t0 = time.time()
im0 = yolov5_detector(im0)
print(f'Done. ({time.time() - t0:.3f}s)')
# print(time.time() - t0)
cv2.imshow("123456", im0)
cv2.waitKey(1) # 1 millisecond
yolov5-python调用精简源码
最新推荐文章于 2024-03-03 16:21:51 发布