有同学后台私信我,想用YOLOv5实现目标的分类计数,因此本文将在之前目标计数博客的基础上添加一些代码,实现分类计数。阅读本文前请先看那篇博客,链接如下:
YOLOv5实现目标计数_Albert_yeager的博客
1. 分类实现
以coco数据集为例,其类别如下(共80类)。注意,每个类别都对应着一个序号,如: 'person' 序号为0, 'bicycle' 序号为1, 'car' 序号为2...这个在之后的调用中会用到。
找到之前写的的计数模块(详见之前的博客),将其替换为下面的代码,即可实现分类计数功能,下面我将进行详细的讲解。
# Write results+计数# count=0person_count = 0tie_count = 0for *xyxy, conf, cls in reversed(det): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format with open(txt_path + '.txt', 'a') as f: f.write(('%g ' * len(line)).rstrip() % line + '\n') if save_img or view_img: # Add bbox to image #c = int(cls)# integer class分类数 #label = '%s %.2f num: %d' % (names[int(cls)], conf, person_count) label = f'{names[int(cls)]} {conf:.2f}' plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) ##########################分类计数########################## if int(cls) == 0: person_count += 1 if int(cls) == 27: tie_count += 1 # count = count+1
添加的主要代码为1、2、4,其中1、2是初始化两个类别的个数,这里我选择人('person')和领带('tie')作为两个计数的类(可以根据需求添加自己的类)。3是显示标签的格式,下面那一行是官方的,你也可以改成自己喜欢的样子。
重点来了!4中的代码是两个判断,int(cls)表示类别的序号。在coco的类别中,'person' 的序号为0,因此当 int(cls) == 0 时也就是当识别到人时,人的计数器 person_count+1;'tie' 的序号为27,因此当 int(cls) == 27 时也就是当识别到领带时,领带的计数器 tie_count+1。这样就能实现分类计数。
这里用的推断模型是yolov5s.pt,是用coco数据集训练出来的,因此识别的类就是上面展示的80类,序号就是依次从0-79(这是一句废话)。那么如果要对自己的数据集进行分类计数,那么就要用自己训练出来的模型进行推断,序号就是按训练部署时names数组中的序列,从0开始,依次递增。
为了讲得更清楚,我再举一个例子(点赞关注一下呗(>﹏<) ,555~)
下面是我自己训练部署时的数据集文件内容,可以看到我要识别五个目标,分别是HEWSN,那么根据names数组中的序列,‘H’对应的序号为0,E对应1,W对应2,S对应3,N对应4。
如果我想对W和S进行分类计数,那么上面的代码应该改成下面这样,首先定义W_count,S_count两个计数器,然后当int(cls) == 2时W_count+1,int(cls) == 3时S_count+1。当然,推断模型记得改成自己的(更改 '--weights' 参数的默认值)。
# Write results+计数# count=0W_count = 0S_count = 0for *xyxy, conf, cls in reversed(det): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format with open(txt_path + '.txt', 'a') as f: f.write(('%g ' * len(line)).rstrip() % line + '\n') if save_img or view_img: # Add bbox to image #c = int(cls)# integer class分类数 #label = '%s %.2f num: %d' % (names[int(cls)], conf, person_count) label = f'{names[int(cls)]} {conf:.2f}' plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) ##########################分类计数########################## if int(cls) == 2: W_count += 1 if int(cls) == 3: S_count += 1 # count = count+1
2. 图片/视频识别显示计数内容
为了将计数结果显示在图像上,需要使用 cv2.putText() 函数,具体添加方法如下:
在 “if save_img:” 后添加下面这几行代码即可,这里我就只打印人和领带的计数了,大家可以根据自己的需求改。
##############################视频识别显示计数内容####################################text = 'person_num:%d ' % (person_count)cv2.putText(im0, text, (180, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 5)text = 'tie_num:%d ' % (tie_count)cv2.putText(im0, text, (180, 120), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 5)####################################################################################
另外注意cv2.putText()函数的几个参数意义,然后慢慢调参,让打印出来的图片美观就行了。
cv2.putText(im0, text, (40, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 4)# 要绘制的图像(im0)# 要绘制的文本字符串(text)# 文本的位置(x, y),窗口左上角为(0,0)# 要使用的字体类型(font),这里用OpenCV的内嵌字体# 字体大小(font_scale),在此处为1# 字体颜色(font_color),在此处为红色(0, 0, 255)# 字体线宽(thickness),在此处为4
最后实现效果如下(视频识别同理):
3. 实时检测窗口打印计数内容
为了将计数结果显示在实时检测的窗口中,同样需要使用 cv2.putText() 函数,具体的修改方法如下:
在 “if view_img:” 后添加下面这几行代码(这里我就只打印人的计数了,大家可以根据自己的需求改)。
##############################实时检测窗口打印计数内容#################################text = 'person_num:%d ' % (person_count)cv2.putText(im0, text, (180, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 4)####################################################################################
希望这篇文章可以帮助到大家,其他评论区和私信问我问题的同学们也不要急,你们的问题我一直在研究,如果完成了我会第一时间发出来并通知你的(≧∇≦)/
求学路上,你我共勉(๑•̀ㅂ•́)و✧