92 lines
4.0 KiB
Python
92 lines
4.0 KiB
Python
import torch
|
|
torch.set_printoptions(precision=10)
|
|
torch.set_printoptions(threshold=99999999999)
|
|
torch.backends.cudnn.deterministic = True
|
|
from .yolov5_preprocess import *
|
|
from .yolov5_postprocess import *
|
|
from .yolo_v2 import Model as Model_v2
|
|
import onnxruntime
|
|
import time
|
|
import os
|
|
from collections import Counter
|
|
import torch.nn.functional as F
|
|
import random
|
|
from pathlib import Path
|
|
|
|
class Yolov5Runner:
|
|
def __init__(self, model_path, yaml_path, grid20_path, grid40_path, grid80_path, num_classes, imgsz_h, imgsz_w, conf_thres, iou_thres, top_k_num, vanish_point, **kwargs):#is_onnx,
|
|
"""
|
|
inputs :
|
|
model_path : str ,path to model
|
|
"""
|
|
self.model_path = model_path
|
|
self.imgsz_h = imgsz_h
|
|
self.imgsz_w = imgsz_w
|
|
self.conf_thres = conf_thres
|
|
self.iou_thres = iou_thres
|
|
self.top_k_num = top_k_num
|
|
self.vanish_point = vanish_point
|
|
self.num_classes = num_classes
|
|
self.DEVICE = torch.device("cpu")#torch.device('cuda:0')#
|
|
self.grid20 = torch.from_numpy(np.load(grid20_path))
|
|
self.grid40 = torch.from_numpy(np.load(grid40_path))
|
|
self.grid80 = torch.from_numpy(np.load(grid80_path))
|
|
self.grids = [self.grid80, self.grid40, self.grid20]
|
|
if 'onnx' not in model_path:
|
|
self.yolov5_model = Model_v2(yaml_path, nc=num_classes)
|
|
self.yolov5_model.load_state_dict(torch.load(model_path, map_location=self.DEVICE))#,strict=False)
|
|
self.yolov5_model.float().eval()
|
|
self.yolov5_model.to(self.DEVICE)
|
|
self.yolov5_model.eval()
|
|
else:
|
|
#onnxruntime.set_default_logger_severity(0)
|
|
self.sess = onnxruntime.InferenceSession(model_path)
|
|
# self.sess.set_providers(['CUDAExecutionProvider'])
|
|
self.input_name = self.sess.get_inputs()[0].name
|
|
self.onnx_batch_size = self.sess.get_inputs()[0].shape[0]
|
|
self.onnx_img_size_h = self.sess.get_inputs()[0].shape[2]
|
|
self.onnx_img_size_w = self.sess.get_inputs()[0].shape[3]
|
|
|
|
self.anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] #yolov5
|
|
print('self.vanish_point',self.vanish_point)
|
|
self.e2e_coco = kwargs.get('e2e_coco', False)
|
|
|
|
def run(self, img_path):
|
|
"""
|
|
inputs :
|
|
img_path : path of the image
|
|
outputs :
|
|
dets : list
|
|
"""
|
|
self.yolov5_model.eval()
|
|
with torch.no_grad():
|
|
img, im0 = Yolov5_preprocess(img_path, self.DEVICE, self.imgsz_h, self.imgsz_w)
|
|
if next(self.yolov5_model.parameters()).is_cuda:
|
|
img = img.type(torch.cuda.FloatTensor)
|
|
else:
|
|
img = img.type(torch.FloatTensor)
|
|
pred = self.yolov5_model(img, augment=False)
|
|
img_shape, im0_shape = img.shape, im0.shape
|
|
dets = Yolov5_postprocess_sig(pred,img_shape, im0_shape, self.conf_thres, self.iou_thres, self.top_k_num, self.grids, self.num_classes, self.anchors,self.vanish_point)
|
|
return dets
|
|
|
|
def run_onnx(self, img_path):
|
|
"""
|
|
inputs :
|
|
img_path : path of the image
|
|
outputs :
|
|
dets : list
|
|
"""
|
|
with torch.no_grad():
|
|
img, im0 = Yolov5_preprocess(img_path, self.DEVICE, self.imgsz_h, self.imgsz_w)
|
|
np_images = np.array(img.cpu())
|
|
np_images = np_images.astype(np.float32)
|
|
pred_onnx = self.sess.run(None, {self.input_name: np_images })
|
|
img_shape, im0_shape = img.shape, im0.shape
|
|
# print('img_shape',img_shape)
|
|
# print('im0_shape', im0_shape)
|
|
dets_onnx = Yolov5_postprocess_onnx_sig(pred_onnx,img_shape, im0_shape, self.conf_thres, self.iou_thres, self.top_k_num, self.grids, self.num_classes, self.anchors,self.vanish_point, self.e2e_coco)
|
|
return dets_onnx
|
|
|
|
|