import sys import os sys.path.insert(0, os.path.abspath(".")) pydir = os.path.dirname(__file__) import warnings import argparse from pathlib import Path import cv2 import torch import torch.nn as nn import numpy as np import val from models.common import Conv from tqdm import tqdm # Disable all warning warnings.filterwarnings("ignore") from models.yolo import DetectionModel from utils.general import non_max_suppression, xywhn2xyxy, xywh2xyxy, scale_boxes, xyxy2xywhn from utils.metrics import ConfusionMatrix, ap_per_class import time import onnxruntime from trt_utils.trt import TrtModel import pycuda.driver as cuda from utils.dataloaders import LoadImagesAndLabels from torch.utils.data import DataLoader names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'} def draw_boxes(model, image_path, mtype): device = torch.device("cuda") image = cv2.imread(image_path) w, h, c = image.shape image_input = cv2.resize(image, (640, 640)) image_input = image_input[:, :, ::-1].transpose(2, 1, 0).astype(np.float32)[np.newaxis, ...] / 255. if mtype == 'ori' or mtype == 'qat': image_input = torch.from_numpy(image_input).to(device) pred = model(image_input) if mtype == "trt": pred = pred[-1].reshape(1, -1, 85) if mtype == "trt" or mtype == "onnx": preds = non_max_suppression(torch.from_numpy(pred).to(torch.device("cuda")), conf_thres=0.1, iou_thres=0.65, max_det=1000, agnostic=False) else: preds = non_max_suppression(pred, conf_thres=0.1, iou_thres=0.65, max_det=1000, agnostic=False) for bboxes in preds: for bbox in bboxes: bbox = xyxy2xywhn(bbox) bbox = xywhn2xyxy(bbox, w, h) bbox = bbox.cpu().numpy() y1,x1,y2,x2,conf,cid = bbox cv2.rectangle(image, (int(x1),int(y1)), (int(x2),int(y2)), (0, 255, 0), 2) label = f'Class: {names[int(cid)]}, Confidence: {conf:.2f}' cv2.putText(image, label, (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA) cv2.imwrite(f"{mtype}.png", image) del model torch.cuda.empty_cache() def compute_metric(model, dataloader, mtype): device = torch.device("cuda") iouv = torch.linspace(0.5, 0.95, 10, device=device) niou = iouv.numel() stats = [] seen = 0 total_model_time = 0. for i in range(2): desc = "warmup" if i == 1: desc = "val" progress_bar = tqdm(dataloader, total=len(dataloader), desc=desc) for data, targets, paths, shapes in progress_bar: data = data.float() / 255. nb, _, height, width = data.shape if mtype == 'ori' or mtype == 'qat': data, targets = data.to(device), targets.to(device) elif mtype == 'onnx' or mtype == 'trt': data, targets = data.numpy(), targets.to(device) else: raise NotImplemented start_time = time.time() pred = model(data) end_time = time.time() if i == 0: continue total_model_time += end_time - start_time if mtype == 'trt': pred = torch.from_numpy(pred[-1].reshape(1, -1, 85)).to(device) elif mtype == "onnx": pred = torch.from_numpy(pred).to(device) else: pass preds = non_max_suppression(pred, conf_thres=0.001, iou_thres=0.6, max_det=300, multi_label=True, agnostic=False) targets[:, 2:] *= torch.tensor((width, height, width, height), device=device) for si, pred in enumerate(preds): seen += 1 labels = targets[targets[:, 0] == si, 1:] nl, npr = labels.shape[0], pred.shape[0] # number of labels, predictions path, shape = Path(paths[si]), shapes[si][0] correct = torch.zeros(npr, niou, dtype=torch.bool, device=device) # init if npr == 0: if nl: stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0])) continue # Predictions predn = pred.clone() scale_boxes(data[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred # Evaluate if nl: tbox = xywh2xyxy(labels[:, 1:5]) # target boxes scale_boxes(data[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels labelsn = torch.cat((labels[:, 0:1], tbox), 1).to(device) # native-space labels correct = val.process_batch(predn, labelsn, iouv) stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0])) # (correct, conf, pcls, tcls) # Compute metrics stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)] # to numpy if len(stats) and stats[0].any(): tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, names=names) ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95 mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() s = ("%22s" + "%11s" * 6) % ("Class", "Images", "P", "R", "mAP50", "mAP50-95", "Time") pf = "%22s" + "%11i" * 1 + "%11.3g" * 5 # print format print(s) print(pf % ("all", seen, mp, mr, map50, map, total_model_time / seen)) class ONNX: def __init__(self, onnx_path, device): sess_options = onnxruntime.SessionOptions() if onnxruntime.get_device() == "GPU": providers = ['CUDAExecutionProvider'] else: providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED self.session = onnxruntime.InferenceSession(onnx_path, sess_options, providers=providers, provider_options=[{"device_id": 0}]*len(providers)) self.input_name = self.get_input_name() self.output_name = self.get_output_name() def get_input_name(self): input_name = [] for node in self.session.get_inputs(): input_name.append(node.name) return input_name def get_output_name(self): output_name = [] for node in self.session.get_outputs(): output_name.append(node.name) return output_name def get_input_feed(self, image): input_feed = {} for name in self.input_name: input_feed[name] = image return input_feed def inference(self, img): input_feed = self.get_input_feed(img) pred = self.session.run(None, input_feed)[0] return pred class TorchModel: def __init__(self, weight, device): self.device = device self.model = self.load_model(weight) def load_model(self, weight): # def load_yolov5_model(model: str, weight, device) -> Model: if 'yolov5l' in weight: cfg = "models/yolov5l.yaml" elif 'yolov5m' in weight: cfg = "models/yolov5m.yaml" elif 'yolov5n' in weight: cfg = "models/yolov5n.yaml" elif 'yolov5s' in weight: cfg = "models/yolov5s.yaml" elif "yolov5x" in weight: cfg = "models/yolov5x.yaml" else: raise NotImplementedError("Only support yolov5[l, m, n, s, x]") model = DetectionModel(cfg=cfg).to(self.device) weight = torch.load(weight, map_location="cpu")["model"].state_dict() model.load_state_dict(weight,strict=False) for m in model.modules(): if type(m) is nn.Upsample: m.recompute_scale_factor = None # torch 1.11.0 compatibility elif type(m) is Conv: m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility model.float() model.eval() with torch.no_grad(): model.fuse() return model @torch.no_grad() def inference(self, img): return self.model(img) def main_eval(args): dataloader = DataLoader(LoadImagesAndLabels(args.data_list, img_size=640, augment=False, rect=False, cache_images=True, single_cls=False), shuffle=False, batch_size=1, collate_fn=LoadImagesAndLabels.collate_fn) if args.mtype == "ori": # # ================ 原始模型 ======================= device = torch.device(f"cuda") model = TorchModel(args.weight, device).inference elif args.mtype=='qat': # =============== qat 模型 ========================== device = torch.device(f"cuda") model = torch.load(args.weight, map_location="cpu")['model'] model.to(device) elif args.mtype=='onnx': # =============== onnx 模型 ====================== device = torch.device(f"cuda") model = ONNX(args.weight, 0).inference # ================ trt 模型 ====================== elif args.mtype=='trt': cuda.init() model = TrtModel(args.weight) compute_metric(model, dataloader, args.mtype) del model torch.cuda.empty_cache() def main_draw(args): if args.mtype == "ori": # # ================ 原始模型 ======================= device = torch.device(f"cuda") model = TorchModel(args.weight, device).inference elif args.mtype=='qat': # =============== qat 模型 ========================== device = torch.device(f"cuda") model = torch.load(args.weight, map_location="cpu")['model'] model.to(device) elif args.mtype=='onnx': # =============== onnx 模型 ====================== device = torch.device(f"cuda") model = ONNX(args.weight, 0).inference # ================ trt 模型 ====================== elif args.mtype=='trt': cuda.init() model = TrtModel(args.weight) draw_boxes(model, args.image_path, args.mtype) del model torch.cuda.empty_cache() if __name__ == "__main__": parser = argparse.ArgumentParser() subps = parser.add_subparsers(dest="cmd") eval_m = subps.add_parser("eval") eval_m.add_argument("--data_list", default="/home/temp/coco2017/val2017.txt", type=str, help="测试数据列表文件路径") eval_m.add_argument("--weight", required=True, type=str, help="权重路径") eval_m.add_argument("--mtype", required=True, type=str, choices=['onnx', 'qat', 'trt', 'ori'], help="模型类型，ori原始模型，qat带有量化节点的模型，onnx模型，tensorrt模型.") draw = subps.add_parser("draw") draw.add_argument("--weight", required=True, type=str, help="输入trt权重路径") draw.add_argument("--image_path", required=True, type=str, help="待检测图片") draw.add_argument("--mtype", required=True, type=str) args = parser.parse_args() if args.cmd == "eval": print(args) main_eval(args) elif args.cmd == "draw": print(args) main_draw(args)