【目标检测】yolov5仅识别人_基于yolo5只检测人

1. 准备

https://blog.csdn.net/qq_44747572/article/details/127604416?spm=1001.2014.3001.5501

2. 相关文件

3. 推理

# coding:gbk
# coding:utf-8
import cv2.cv2 as cv2
import numpy as np
import onnxruntime
import torch
import torchvision
import time
import random
from utils.general import non_max_suppression
import os


class YOLOV5_ONNX(object):
	def __init__(self,onnx_path):
		'''初始化onnx'''
		self.onnx_session=onnxruntime.InferenceSession(onnx_path)
		print(onnxruntime.get_device())
		self.input_name=self.get_input_name()
		self.output_name=self.get_output_name()
		self.classes=['person', 'car', 'special_person', 'truck']
	def get_input_name(self):
		'''获取输入节点名称'''
		input_name=[]
		for node in self.onnx_session.get_inputs():
			input_name.append(node.name)

		return input_name


	def get_output_name(self):
		'''获取输出节点名称'''
		output_name=[]
		for node in self.onnx_session.get_outputs():
			output_name.append(node.name)

		return output_name

	def get_input_feed(self,image_tensor):
		'''获取输入tensor'''
		input_feed={}
		for name in self.input_name:
			input_feed[name]=image_tensor

		return input_feed

	def letterbox(self,img, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True,
				  stride=32):
		'''图片归一化'''
		# Resize and pad image while meeting stride-multiple constraints
		shape = img.shape[:2]  # current shape [height, width]
		if isinstance(new_shape, int):
			new_shape = (new_shape, new_shape)

		# Scale ratio (new / old)
		r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
		if not scaleup:  # only scale down, do not scale up (for better test mAP)
			r = min(r, 1.0)

		# Compute padding
		ratio = r, r  # width, height ratios

		new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
		dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding

		if auto:  # minimum rectangle
			dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
		elif scaleFill:  # stretch
			dw, dh = 0.0, 0.0
			new_unpad = (new_shape[1], new_shape[0])
			ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

		dw /= 2  # divide padding into 2 sides
		dh /= 2

		if shape[::-1] != new_unpad:  # resize
			img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)

		top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
		left, right = int(round(dw - 0.1)), int(round(dw + 0.1))

		img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
		return img, ratio, (dw, dh)

	def xywh2xyxy(self,x):
		# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
		y = np.copy(x)

		y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
		y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
		y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
		y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y

		return y

	def nms(self,prediction, conf_thres=0.1, iou_thres=0.6, agnostic=False):
		if prediction.dtype is torch.float16:
			prediction = prediction.float()  # to FP32
		xc = prediction[..., 4] > conf_thres  # candidates
		min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
		max_det = 300  # maximum number of detections per image
		output = [None] * prediction.shape[0]
		for xi, x in enumerate(prediction):  # image index, image inference
			x = x[xc[xi]]  # confidence
			if not x.shape[0]:
				continue

			x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
			box = self.xywh2xyxy(x[:, :4])

			conf, j = x[:, 5:].max(1, keepdim=True)
			x = torch.cat((torch.tensor(box), conf, j.float()), 1)[conf.view(-1) > conf_thres]
			n = x.shape[0]  # number of boxes
			if not n:
				continue
			c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
			boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
			i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
			if i.shape[0] > max_det:  # limit detections
				i = i[:max_det]
			output[xi] = x[i]

		return output

	def clip_coords(self,boxes, img_shape):
		'''查看是否越界'''
		# Clip bounding xyxy bounding boxes to image shape (height, width)
		boxes[:, 0].clamp_(0, img_shape[1])  # x1
		boxes[:, 1].clamp_(0, img_shape[0])  # y1
		boxes[:, 2].clamp_(0, img_shape[1])  # x2
		boxes[:, 3].clamp_(0, img_shape[0])  # y2

	def scale_coords(self,img1_shape, coords, img0_shape, ratio_pad=None):
		'''
		坐标对应到原始图像上，反操作：减去pad，除以最小缩放比例
		:param img1_shape: 输入尺寸
		:param coords: 输入坐标
		:param img0_shape: 映射的尺寸
		:param ratio_pad:
		:return:
		'''

		# Rescale coords (xyxy) from img1_shape to img0_shape
		if ratio_pad is None:  # calculate from img0_shape
			gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new,计算缩放比率
			pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (
						img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding ，计算扩充的尺寸
		else:
			gain = ratio_pad[0][0]
			pad = ratio_pad[1]

		coords[:, [0, 2]] -= pad[0]  # x padding，减去x方向上的扩充
		coords[:, [1, 3]] -= pad[1]  # y padding，减去y方向上的扩充
		coords[:, :4] /= gain  # 将box坐标对应到原始图像上
		self.clip_coords(coords, img0_shape)  # 边界检查
		return coords

	def sigmoid(self,x):
		return 1 / (1 + np.exp(-x))



	def infer(self,img_path):
		'''执行前向操作预测输出'''
		# 超参数设置
		img_size=(640,640) #图片缩放大小
		# 读取图片
		src_img=cv2.imread(img_path)
		start=time.time()
		src_size=src_img.shape[:2]

		# 图片填充并归一化
		img=self.letterbox(src_img,img_size,stride=32)[0]

		# Convert
		img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
		img = np.ascontiguousarray(img)


		# 归一化
		img=img.astype(dtype=np.float32)
		img/=255.0

		# # BGR to RGB
		# img = img[:, :, ::-1].transpose(2, 0, 1)
		# img = np.ascontiguousarray(img)

		# 维度扩张
		img=np.expand_dims(img,axis=0)
		# print('img resuming: ',time.time()-start)
		# 前向推理
		# start=time.time()
		input_feed=self.get_input_feed(img)
		# ort_inputs = {self.onnx_session.get_inputs()[0].name: input_feed[None].numpy()}
		pred = torch.tensor(self.onnx_session.run(None, input_feed)[0])
		results = non_max_suppression(pred, 0.5,0.5)
		# print('onnx resuming: ',time.time()-start)
		# pred=self.onnx_session.run(output_names=self.output_name,input_feed=input_feed)

		#映射到原始图像
		img_shape=img.shape[2:]
		# print(img_size)
		for det in results:  # detections per image
			if det is not None and len(det):
				det[:, :4] = self.scale_coords(img_shape, det[:, :4],src_size).round()
		# print(time.time()-start)
		if det is not None and len(det):
			self.draw(src_img, det)


	def plot_one_box(self,x, img, color=None, label=None, line_thickness=None):
		# Plots one bounding box on image img
		tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
		color = color or [random.randint(0, 255) for _ in range(3)]
		c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
		cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
		if label:
			tf = max(tl - 1, 1)  # font thickness
			t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
			c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
			cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
			cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

	def draw(self,img, boxinfo):
		colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(self.classes))]
		for *xyxy, conf, cls in boxinfo:
			label = '%s %.2f' % ('person', conf)
			if int(cls) == 0:    # 仅仅判断person
				self.plot_one_box(xyxy, img, label=label, color=[0,0,255], line_thickness=1)
		cv2.namedWindow("dst",0)
		cv2.imshow("dst", img)
		cv2.imwrite("res1.jpg",img)
		cv2.waitKey(0)
		return 0


if __name__=="__main__":
	model=YOLOV5_ONNX(onnx_path="./yolov5s.onnx")
	model.infer(img_path="./zidane.jpg")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239