code
import os
import cv2
import argparse
import time
import torch
import numpy as np
from models.common import DetectMultiBackend
from utils.torch_utils import select_device, smart_inference_mode
from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,
letterbox, mixup, random_perspective)
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
from ultralytics.utils.plotting import Annotator, colors, save_one_box
class Yolov5:
def __init__(self, model_path, img_size, conf_thres, iou_thres):
self.img_size=img_size
self.conf_thres=conf_thres
self.iou_thres=iou_thres
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = DetectMultiBackend(model_path, device=self.device , dnn=False, data='data/VOC.yaml', fp16=False)
self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
def preprocess(self, img):
im = letterbox(img, self.img_size, stride=self.stride, auto=True)[0] # padded resize
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
im = np.ascontiguousarray(im) # contiguous
return img
def inference(self, img):
im0=img
imShape=img.shape
gn = torch.tensor(img.shape)[[1, 0, 1, 0]]
img = self.preprocess(img)
with torch.no_grad():
im = torch.from_numpy(img).to(self.device)
im=im.float()
im = im.unsqueeze(0)
im=im.permute(0, 3, 1, 2)
im /= 255
print(im.shape)
pred = self.model(im, augment=False, visualize=False)
conf_thres=0.25 # confidence threshold
iou_thres=0.45 # NMS IOU threshold
max_det=1000 # maximum detections per image)
classes=[0]
print(conf_thres)
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, False, max_det=max_det)
for i, det in enumerate(pred):
annotator = Annotator(im0, line_width=3, example=str(self.names))
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], imShape).round()
# Write results
for *xyxy, conf, cls in reversed(det):
c = int(cls) # integer class
confidence = float(conf)
confidence_str = f'{confidence:.2f}'
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()
annotator.box_label(xyxy,'C', color=colors(c, True))
im0 = annotator.result()
cv2.imwrite('res.jpg',im0)
return ''