evaluate.py

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)