Tutorial_Python.md 6.31 KB
Newer Older
liucong's avatar
liucong committed
1
2
# YOLOV5检测器

3
本份文档主要介绍如何基于MIGraphX构建YOLOV5的动态shape推理Python示例,根据文档描述可以了解怎样运行该Python示例得到YOLOV5的目标检测结果。
liucong's avatar
liucong committed
4
5
6
7
8
9
10
11
12
13
14
15
16

## 模型简介

YOLOV5是一种单阶段目标检测算法,该算法在YOLOV4的基础上添加了一些新的改进思路,使其速度与精度都得到了极大的性能提升。具体包括:输入端的Mosaic数据增强、自适应锚框计算、自适应图像缩放操作;主干网络的Focus结构与CSP结构;Neck端的FPN+PAN结构;输出端的损失函数GIOU_Loss以及预测框筛选的DIOU_nms。网络结构如图所示。

<img src=./YOLOV5_01.jpg style="zoom:100%;" align=middle>

## 预处理

待检测图像输入模型进行检测之前需要进行预处理,主要包括调整输入的尺寸,归一化等操作。

1. 转换数据排布为NCHW
2. 归一化[0.0, 1.0]
17
3. 调整输入数据的尺寸
liucong's avatar
liucong committed
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

```
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))
    # 维度转换HWC->CHW
    input_img = input_img.transpose(2, 0, 1)
    # 维度拓展,增加batch维度
    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
```

37
38
39
40
41
其中模型输入的inputWidth、inputHeight为每次动态输入shape。

## 推理

执行YOLOV5动态输入推理,首先需要对YOLOV5的动态模型进行解析、编译,与静态推理不同的是,动态shape推理需要设置模型输入的最大shape,本示例设为[1,3,800,800]。
liucong's avatar
liucong committed
42
43
44
45
46
47
48
49
50
51
52
53

```
class YOLOv5:
    def __init__(self, path, 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.classNames = list(map(lambda x: x.strip(), open('../Resource/Models/coco.names', 'r').readlines()))

        # 解析推理模型
54
55
        maxInput={"images":[1,3,800,800]}
        self.model = migraphx.parse_onnx(path, map_input_dims=maxInput)
liucong's avatar
liucong committed
56
57
58
59
60
61

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

        # 获取模型的输入尺寸
        inputShape = self.model.get_parameter_shapes()[self.inputName].lens()
62
63
64
65
66
67
68
        print("inputName:{0} \ninputMaxShape:{1}".format(self.inputName, inputShape))
        
        # 模型编译
        self.model.compile(t=migraphx.get_target("gpu"), device_id=0)  # device_id: 设置GPU设备,默认为0号设备
        print("Success to compile")
        
        ...
liucong's avatar
liucong committed
69
70
```

71
模型初始化完成之后开始进行推理,对输入数据进行前向计算得到模型的输出result,在detect函数中调用定义的process_output函数对result进行后处理,得到图像中含有物体的anchor坐标信息、类别置信度、类别ID。
liucong's avatar
liucong committed
72
73

```
74
75
76
77
def detect(self, image, input_shape):
    self.inputWidth = input_shape[3]
    self.inputHeight = input_shape[2]
    # 输入图片预处理
liucong's avatar
liucong committed
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    input_img = self.prepare_input(image)

    # 执行推理
    start = time.time()
    result = self.model.run({self.model.get_parameter_names()[0]: input_img})
    print('net forward time: {:.4f}'.format(time.time() - start))
    # 模型输出结果后处理
    boxes, scores, class_ids = self.process_output(result)

    return boxes, scores, class_ids
```

其中对migraphx推理输出result进行后处理,首先需要对生成的anchor根据是否有物体阈值objectThreshold、置信度阈值confThreshold进行筛选,相关过程定义在process_output函数中。获取筛选后的anchor的坐标信息之后,需要将坐标映射到原图中的位置,相关过程定义在rescale_boxes函数中。

```
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]

119
120
121
122
123
124
125
126
127
128
129
130
131
def extract_boxes(self, predictions):
    # 获取anchor的坐标信息
    boxes = predictions[:, :4]

    # 将anchor的坐标信息映射到输入image
    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_

liucong's avatar
liucong committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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
```

根据获取的detect函数输出的boxes、scores、class_ids信息在原图进行结果可视化,包括绘制图像中检测到的物体位置、类别和置信度分数,得到最终的YOLOV5目标检测结果输出。

```
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} {int(score * 100)}%'
        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
```