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

        有同学后台私信我，想用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)
####################################################################################

        

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

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