yolox_infer_migraphx.py

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

def nms(boxes, scores, nms_thr):
    """Single class NMS implemented in Numpy."""
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= nms_thr)[0]
        order = order[inds + 1]

    return keep

def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
        """Multiclass NMS implemented in Numpy"""
        if class_agnostic:
            nms_method = multiclass_nms_class_agnostic
        else:
            nms_method = multiclass_nms_class_aware
        return nms_method(boxes, scores, nms_thr, score_thr)

def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
    """Multiclass NMS implemented in Numpy. Class-aware version."""
    final_dets = []
    num_classes = scores.shape[1]
    for cls_ind in range(num_classes):
        cls_scores = scores[:, cls_ind]
        valid_score_mask = cls_scores > score_thr
        if valid_score_mask.sum() == 0:
            continue
        else:
            valid_scores = cls_scores[valid_score_mask]
            valid_boxes = boxes[valid_score_mask]
            keep = nms(valid_boxes, valid_scores, nms_thr)
            if len(keep) > 0:
                cls_inds = np.ones((len(keep), 1)) * cls_ind
                dets = np.concatenate(
                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
                )
                final_dets.append(dets)
    if len(final_dets) == 0:
        return None
    return np.concatenate(final_dets, 0)

def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
    """Multiclass NMS implemented in Numpy. Class-agnostic version."""
    cls_inds = scores.argmax(1)
    cls_scores = scores[np.arange(len(cls_inds)), cls_inds]

    valid_score_mask = cls_scores > score_thr
    if valid_score_mask.sum() == 0:
        return None
    valid_scores = cls_scores[valid_score_mask]
    valid_boxes = boxes[valid_score_mask]
    valid_cls_inds = cls_inds[valid_score_mask]
    keep = nms(valid_boxes, valid_scores, nms_thr)
    if keep:
        dets = np.concatenate(
            [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], 1
        )
    return dets

_COLORS = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.300, 0.300, 0.300,
        0.600, 0.600, 0.600,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.749, 0.749, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 1.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.286, 0.286, 0.286,
        0.429, 0.429, 0.429,
        0.571, 0.571, 0.571,
        0.714, 0.714, 0.714,
        0.857, 0.857, 0.857,
        0.000, 0.447, 0.741,
        0.314, 0.717, 0.741,
        0.50, 0.5, 0
    ]
).astype(np.float32).reshape(-1, 3)

def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):
    for i in range(len(boxes)):
        box = boxes[i]
        cls_id = int(cls_ids[i])
        score = scores[i]
        if score < conf:
            continue
        x0 = int(box[0])
        y0 = int(box[1])
        x1 = int(box[2])
        y1 = int(box[3])

        color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
        text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)
        txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
        font = cv2.FONT_HERSHEY_SIMPLEX

        txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
        cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)

        txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()
        cv2.rectangle(
            img,
            (x0, y0 + 1),
            (x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])),
            txt_bk_color,
            -1
        )
        cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)

    return img


class YOLOX:
    """YOLOX object detection model class for handling inference and visualization."""

    def __init__(self, model_path, dynamic=False, conf_thres=0.5, iou_thres=0.5):
        """
        Initializes an instance of the YOLOX class.

        Args:
            model_path: Path to the ONNX model.
            dynamic: whether use dynamic inference.
            conf_thres: Confidence threshold for filtering detections.
            iou_thres: IoU (Intersection over Union) threshold for non-maximum suppression.
        """
        self.confThreshold = conf_thres
        self.nmsThreshold = iou_thres
        self.isDynamic = dynamic
        # 获取模型检测的类别信息
        self.classNames = list(map(lambda x: x.strip(), open('/home/yolox_migraphx/Resource/Models/coco.names', 'r').readlines()))

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

            # 获取模型输入/输出节点信息
            print("inputs:")
            inputs = self.model.get_inputs()
            for key,value in inputs.items():
                print("{}:{}".format(key,value))
            
            print("outputs:")
            outputs = self.model.get_outputs()
            for key,value in outputs.items():
                print("{}:{}".format(key,value))

            # 获取模型的输入name
            self.inputName = "images"
            
            # 获取模型的输入尺寸
            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(model_path) 
            # 获取模型输入/输出节点信息
            print("inputs:")
            inputs = self.model.get_inputs()
            for key,value in inputs.items():
                print("{}:{}".format(key,value))
            
            print("outputs:")
            outputs = self.model.get_outputs()
            for key,value in outputs.items():
                print("{}:{}".format(key,value))

            # 获取模型的输入name
            self.inputName = "images"

            # 获取模型的输入尺寸
            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号设备
        print("Success to compile")

        # Generate a color palette for the classes
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classNames), 3))

    # def draw_detections(self, img, box, score, class_id):
    #     """
    #     Draws bounding boxes and labels on the input image based on the detected objects.

    #     Args:
    #         img: The input image to draw detections on.
    #         box: Detected bounding box.
    #         score: Corresponding detection score.
    #         class_id: Class ID for the detected object.

    #     Returns:
    #         None
    #     """

    #     # Extract the coordinates of the bounding box
    #     x1, y1, w, h = box

    #     # Retrieve the color for the class ID
    #     color = self.color_palette[class_id]

    #     # Draw the bounding box on the image
    #     cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)

    #     # Create the label text with class name and score
    #     label = f'{self.classNames[class_id]}: {score:.2f}'

    #     # Calculate the dimensions of the label text
    #     (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

    #     # Calculate the position of the label text
    #     label_x = x1
    #     label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10

    #     # Draw a filled rectangle as the background for the label text
    #     cv2.rectangle(img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color,
    #                   cv2.FILLED)

    #     # Draw the label text on the image
    #     cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)

    # def preprocess(self, image):
    #     """
    #     Preprocesses the input image before performing inference.

    #     Returns:
    #         image_data: Preprocessed image data ready for inference.
    #     """
    #     # Read the input image using OpenCV
    #     # self.img = cv2.imread(self.input_image)
    #     self.img = image

    #     # Get the height and width of the input image
    #     self.img_height, self.img_width = self.img.shape[:2]

    #     # Convert the image color space from BGR to RGB
    #     img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)

    #     # Resize the image to match the input shape
    #     img = cv2.resize(img, (self.inputWidth, self.inputHeight))

    #     # Normalize the image data by dividing it by 255.0
    #     image_data = np.array(img) / 255.0

    #     # Transpose the image to have the channel dimension as the first dimension
    #     image_data = np.transpose(image_data, (2, 0, 1))  # Channel first

    #     # Expand the dimensions of the image data to match the expected input shape
    #     image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

    #     # Make array memery contiguous
    #     image_data = np.ascontiguousarray(image_data)

    #     # Return the preprocessed image data
    #     return image_data

    # def postprocess(self, input_image, output):
    #     """
    #     Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs.

    #     Args:
    #         input_image (numpy.ndarray): The input image.
    #         output (numpy.ndarray): The output of the model.

    #     Returns:
    #         numpy.ndarray: The input image with detections drawn on it.
    #     """

    #     # Transpose and squeeze the output to match the expected shape
    #     outputs = np.transpose(np.squeeze(output[0]))

    #     # Get the number of rows in the outputs array
    #     rows = outputs.shape[0]

    #     # Lists to store the bounding boxes, scores, and class IDs of the detections
    #     boxes = []
    #     scores = []
    #     class_ids = []

    #     # Calculate the scaling factors for the bounding box coordinates
    #     x_factor = self.img_width / self.inputWidth
    #     y_factor = self.img_height / self.inputHeight

    #     # Iterate over each row in the outputs array
    #     for i in range(rows):
    #         # Extract the class scores from the current row
    #         classes_scores = outputs[i][4:]

    #         # Find the maximum score among the class scores
    #         max_score = np.amax(classes_scores)

    #         # If the maximum score is above the confidence threshold
    #         if max_score >= self.confThreshold:
    #             # Get the class ID with the highest score
    #             class_id = np.argmax(classes_scores)

    #             # Extract the bounding box coordinates from the current row
    #             x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

    #             # Calculate the scaled coordinates of the bounding box
    #             left = int((x - w / 2) * x_factor)
    #             top = int((y - h / 2) * y_factor)
    #             width = int(w * x_factor)
    #             height = int(h * y_factor)

    #             # Add the class ID, score, and box coordinates to the respective lists
    #             class_ids.append(class_id)
    #             scores.append(max_score)
    #             boxes.append([left, top, width, height])

    #     # Apply non-maximum suppression to filter out overlapping bounding boxes
    #     indices = cv2.dnn.NMSBoxes(boxes, scores, self.confThreshold, self.nmsThreshold)

    #     # Iterate over the selected indices after non-maximum suppression
    #     for i in indices:
    #         # Get the box, score, and class ID corresponding to the index
    #         box = boxes[i]
    #         score = scores[i]
    #         class_id = class_ids[i]

    #         # Draw the detection on the input image
    #         self.draw_detections(input_image, box, score, class_id)

    #     # Return the modified input image
    #     return input_image

    def preproc(self, img, input_size, swap=(2, 0, 1)):
        if len(img.shape) == 3:
            padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
        else:
            padded_img = np.ones(input_size, dtype=np.uint8) * 114

        r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * r), int(img.shape[0] * r)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

        padded_img = padded_img.transpose(swap)
        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
        padded_img = np.expand_dims(padded_img, axis=0)
        return padded_img, r

    def demo_postprocess(self, outputs, img_size, p6=False):
        grids = []
        expanded_strides = []
        strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]

        hsizes = [img_size[0] // stride for stride in strides]
        wsizes = [img_size[1] // stride for stride in strides]

        for hsize, wsize, stride in zip(hsizes, wsizes, strides):
            xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
            grids.append(grid)
            shape = grid.shape[:2]
            expanded_strides.append(np.full((*shape, 1), stride))

        grids = np.concatenate(grids, 1)
        expanded_strides = np.concatenate(expanded_strides, 1)
        outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
        outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides

        return outputs
    
    def detect(self, image, output_dir, image_path, input_shape=None):
        # if(self.isDynamic):
        #     self.inputWidth = input_shape[3]
        #     self.inputHeight = input_shape[2]
        # 输入图片预处理
        # input_img = self.preprocess(image)
        img, ratio = self.preproc(image, input_shape)
        flat_predictions = np.array(img).ravel()
        for i in range(min(100, len(flat_predictions))):
            print(flat_predictions[i])
        # 执行推理
        start = time.time()
        result = self.model.run({self.inputName: img})
        print('net forward time: {:.4f}'.format(time.time() - start))
        # 模型输出结果后处理
        # dstimg = self.postprocess(image, result)
        predictions = self.demo_postprocess(np.array(result[0]), input_shape)[0]
        # flat_predictions = np.array(result[0]).ravel()
        # for i in range(min(100, len(flat_predictions))):
        #     print(flat_predictions[i])

        boxes = predictions[:, :4]
        scores = predictions[:, 4:5] * predictions[:, 5:]
        print("max(predictions[:, 4:5]):{}".format(np.amax(predictions[:, 4:5])))
        print("max(predictions[:, 5:]):{}".format(np.amax(predictions[:, 5:])))
        print("max(scores):{}".format(np.amax(scores)))
        boxes_xyxy = np.ones_like(boxes)
        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
        boxes_xyxy /= ratio
        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=self.nmsThreshold, score_thr=0.1)
        if dets is not None:
            final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
            origin_img = vis(image, final_boxes, final_scores, final_cls_inds,
                            conf=self.confThreshold, class_names=COCO_CLASSES)

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        output_path = os.path.join(output_dir, os.path.basename(image_path))
        cv2.imwrite(output_path, origin_img)

        return origin_img


def read_images(image_path):
    image_lists = []
    image_name_lists = []
    
    for image_name in os.listdir(image_path):
        image = cv2.imread(image_path +"/" + image_name, 1)
        image_name_lists.append(image_path +"/" + image_name)
        image_lists.append(image)
        
    return image_lists, image_name_lists

def yoloX_Static(imgpath, modelpath, confThreshold, nmsThreshold, output_dir, input_shape):
    yoloX_detector = YOLOX(modelpath, False, conf_thres=confThreshold,
                             iou_thres=nmsThreshold)
    srcimg = cv2.imread(imgpath, 1)

    dstimg = yoloX_detector.detect(srcimg, output_dir, imgpath, input_shape)

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


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

if __name__ == '__main__':
    # Create an argument parser to handle command-line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgPath', type=str, default='/home/yolox_migraphx/Resource/Images/image_test.jpg', help="image path")
    parser.add_argument('--imgFolderPath', type=str, default='/home/yolox_migraphx/Resource/Images/DynamicPics', help="image folder path")
    parser.add_argument('--staticModelPath', type=str, default='/home/yolox_migraphx/Resource/Models/yolox_s.onnx', help="static onnx filepath")
    parser.add_argument('--dynamicModelPath', type=str, default='/home/yolox_migraphx/Resource/Models/yolox_s_dynamic.onnx', help="dynamic onnx filepath")
    parser.add_argument('--confThreshold', default=0.5, 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 dynamic inference")
    parser.add_argument(
        "-o",
        "--output_dir",
        type=str,
        default='demo_output',
        help="Path to your output directory.",
    )
    parser.add_argument(
        "--input_shape",
        type=str,
        default="640,640",
        help="Specify an input shape for inference.",
    )
    args = parser.parse_args()
    input_shape = [int(dim) for dim in args.input_shape.split(",")]
    
    # 静态推理
    if args.staticInfer:
        yoloX_Static(args.imgPath, args.staticModelPath, args.confThreshold, args.nmsThreshold, args.output_dir, input_shape)
    # 动态推理
    if args.dynamicInfer:
        yoloX_dynamic(args.imgFolderPath, args.dynamicModelPath, args.confThreshold, args.nmsThreshold, args.output_dir, input_shape)