YOLOv5实现目标分类计数并显示在图像上

        有同学后台私信我,想用YOLOv5实现目标的分类计数,因此本文将在之前目标计数博客的基础上添加一些代码,实现分类计数。阅读本文前请先看那篇博客,链接如下:

YOLOv5实现目标计数_Albert_yeager的博客

1. 分类实现

        以coco数据集为例,其类别如下(共80类)。注意,每个类别都对应着一个序号,如: 'person' 序号为0,  'bicycle' 序号为1,  'car'  序号为2...这个在之后的调用中会用到。

         找到之前写的的计数模块(详见之前的博客),将其替换为下面的代码,即可实现分类计数功能,下面我将进行详细的讲解。

# Write results+计数
# count=0
person_count = 0
tie_count = 0
for *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=0
W_count = 0
S_count = 0
for *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)
####################################################################################

        

        希望这篇文章可以帮助到大家,其他评论区和私信问我问题的同学们也不要急,你们的问题我一直在研究,如果完成了我会第一时间发出来并通知你的(≧∇≦)/

求学路上,你我共勉(๑•̀ㅂ•́)و✧

猜你喜欢

转载自blog.csdn.net/Albert_yeager/article/details/130694180