YoloV5_infer_migraphx.py

# -*- coding: utf-8 -*-
import cv2
import numpy as np
import os
import argparse
import time
import migraphx


class YOLOv5:
    def __init__(self, path, dynamic=False, obj_thres=0.5, conf_thres=0.25, iou_thres=0.5):
        self.objectThreshold = obj_thres
        self.confThreshold = conf_thres
        self.nmsThreshold = iou_thres
        self.isDynamic = dynamic
        # 获取模型检测的类别信息
        self.classNames = list(map(lambda x: x.strip(), open('../Resource/Models/coco.names', 'r').readlines()))

        # 解析推理模型
        if self.isDynamic:
            maxInput={"images":[1,3,800,800]}
            self.model = migraphx.parse_onnx(path, map_input_dims=maxInput)

            # 获取模型输入/输出节点信息
            inputs = self.model.get_inputs()
            outputs = self.model.get_outputs()

            # 获取模型的输入name
            self.inputName = self.model.get_parameter_names()[0]
            
            # 获取模型的输入尺寸
            inputShape = inputShape=inputs[self.inputName].lens()
            self.inputHeight = int(inputShape[2])
            self.inputWidth = int(inputShape[3])
            print("inputName:{0} \ninputShape:{1}".format(self.inputName, inputShape))
        else:
            self.model = migraphx.parse_onnx(path) 
            # 获取模型输入/输出节点信息
            inputs = self.model.get_inputs()
            outputs = self.model.get_outputs()

            # 获取模型的输入name
            self.inputName = self.model.get_parameter_names()[0]

            # 获取模型的输入尺寸
            inputShape = inputShape=inputs[self.inputName].lens()
            self.inputHeight = int(inputShape[2])
            self.inputWidth = int(inputShape[3])
            print("inputName:{0} \ninputShape:{1}".format(self.inputName, inputShape))
        
        # 模型编译
        self.model.compile(t=migraphx.get_target("gpu"), device_id=0)  # device_id: 设置GPU设备，默认为0号设备

    def detect(self, image, input_shape=None):
        if(self.isDynamic):
            self.inputWidth = input_shape[3]
            self.inputHeight = input_shape[2]
        # 输入图片预处理
        input_img = self.prepare_input(image)

        # 执行推理
        result = self.model.run({self.inputName: input_img})

        # 模型输出结果后处理
        boxes, scores, class_ids = self.process_output(result)

        return boxes, scores, class_ids

    def prepare_input(self, image):
        self.img_height, self.img_width = image.shape[:2]
        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        input_img = cv2.resize(input_img, (self.inputWidth, self.inputHeight))
        input_img = input_img.transpose(2, 0, 1)
        input_img = np.expand_dims(input_img, 0)
        input_img = np.ascontiguousarray(input_img)
        input_img = input_img.astype(np.float32)
        input_img = input_img / 255

        return input_img

    def process_output(self, output):
        predictions = np.squeeze(output[0])

        # 筛选包含物体的anchor
        obj_conf = predictions[:, 4]
        predictions = predictions[obj_conf > self.objectThreshold]
        obj_conf = obj_conf[obj_conf > self.objectThreshold]

        # 筛选大于置信度阈值的anchor
        predictions[:, 5:] *= obj_conf[:, np.newaxis]
        scores = np.max(predictions[:, 5:], axis=1)
        valid_scores = scores > self.confThreshold
        predictions = predictions[valid_scores]
        scores = scores[valid_scores]

        # 获取最高置信度分数对应的类别ID
        class_ids = np.argmax(predictions[:, 5:], axis=1)

        # 获取每个物体对应的anchor
        boxes = self.extract_boxes(predictions)

        # 执行非极大值抑制消除冗余anchor
        indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.confThreshold, self.nmsThreshold).flatten()

        return boxes[indices], scores[indices], class_ids[indices]

    def extract_boxes(self, predictions):
        boxes = predictions[:, :4]
        boxes = self.rescale_boxes(boxes)
        boxes_ = np.copy(boxes)
        boxes_[..., 0] = boxes[..., 0] - boxes[..., 2] * 0.5
        boxes_[..., 1] = boxes[..., 1] - boxes[..., 3] * 0.5
        return boxes_

    def rescale_boxes(self, boxes):
        # 对anchor尺寸进行变换
        input_shape = np.array([self.inputWidth, self.inputHeight, self.inputWidth, self.inputHeight])
        boxes = np.divide(boxes, input_shape, dtype=np.float32)
        boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
        return boxes

    def draw_detections(self, image, boxes, scores, class_ids):
        for box, score, class_id in zip(boxes, scores, class_ids):
            cx, cy, w, h = box.astype(int)

            # 绘制检测物体框
            cv2.rectangle(image, (cx, cy), (cx + w, cy + h), (0, 255, 255), thickness=2)
            label = self.classNames[class_id]
            label = f'{label} {score:.2f}'
            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.putText(image, label, (cx, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), thickness=2)
        return image
    
def read_images(image_path):
    image_lists = []
    
    for image_name in os.listdir(image_path):
        image = cv2.imread(image_path +"/" + image_name, 1)
        image_lists.append(image)
        
    return image_lists

def yolov5_Static(imgpath, modelpath, objectThreshold, confThreshold, nmsThreshold):
    yolov5_detector = YOLOv5(modelpath, False, obj_thres=objectThreshold, conf_thres=confThreshold,
                             iou_thres=nmsThreshold)
    srcimg = cv2.imread(imgpath, 1)

    boxes, scores, class_ids = yolov5_detector.detect(srcimg)

    dstimg = yolov5_detector.draw_detections(srcimg, boxes, scores, class_ids)

    # 保存检测结果
    cv2.imwrite("./Result.jpg", dstimg)
    print("Success to save result")


def yolov5_dynamic(imgpath, modelpath, objectThreshold, confThreshold, nmsThreshold):
    # 设置动态输入shape
    input_shapes = []
    input_shapes.append([1,3,416,416])
    input_shapes.append([1,3,608,608])
    
    # 读取测试图像
    image_lists = read_images(imgpath)
    
    # 推理
    yolov5_detector = YOLOv5(modelpath, True, obj_thres=objectThreshold, 
                                    conf_thres=confThreshold, iou_thres=nmsThreshold)
    for i, image in enumerate(image_lists):
        print("Start to inference image{}".format(i))
        boxes, scores, class_ids = yolov5_detector.detect(image, input_shapes[i])
        dstimg = yolov5_detector.draw_detections(image, boxes, scores, class_ids)
        
        # 保存检测结果
        result_name = "Result{}.jpg".format(i)
        cv2.imwrite(result_name, dstimg)
    
    print("Success to save results")
    
    

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgPath', type=str, default='../Resource/Images/DynamicPics/image1.jpg', help="image path")
    parser.add_argument('--imgFolderPath', type=str, default='../Resource/Images/DynamicPics', help="image folder path")
    parser.add_argument('--staticModelPath', type=str, default='../Resource/Models/yolov5s.onnx', help="static onnx filepath")
    parser.add_argument('--dynamicModelPath', type=str, default='../Resource/Models/yolov5s_Nx3xNxN.onnx', help="static onnx filepath")
    parser.add_argument('--objectThreshold', default=0.5, type=float, help='class confidence')
    parser.add_argument('--confThreshold', default=0.25, type=float, help='class confidence')
    parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh')
    parser.add_argument("--staticInfer",action="store_true",default=False,help="Performing static inference")
    parser.add_argument("--dynamicInfer",action="store_true",default=False,help="Performing static inference")
    args = parser.parse_args()
    
    # 静态推理
    if args.staticInfer:
        yolov5_Static(args.imgPath, args.staticModelPath, args.objectThreshold, args.confThreshold, args.nmsThreshold)
        
    # 动态推理
    if args.dynamicInfer:
        yolov5_dynamic(args.imgFolderPath, args.dynamicModelPath, args.objectThreshold, args.confThreshold, args.nmsThreshold)