YOLOv8 Track部分类别计数器

基于YOLOv8的计数器

先讲核心计数代码result.py中的plot函数

for d in reversed(pred_boxes):
                c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())

c, conf, id 分别代表类别号、置信度、id号，除第一次检测外，每一次有新的目标出现都会增加id号（id是所有类别拉通排号的），因此使用max_id存储前一次的id，如果检测到id号大于max_id，说明有新的目标出现，并且可能是一次出现多个目标，使用for循环遍历每一次检测到的所有目标，并在id对应的dict键值上＋1。

        # Plot Detect results
        if pred_boxes and show_boxes:
            for d in reversed(pred_boxes):

                c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
                # 第一帧
                if self.max_id == 0:
                    self.max_id = id
                    for i in pred_boxes.cls:
                        self.count_num[names[i.item()]] += 1 
                ##################用于计数################
                elif id is not None and id > self.max_id:
                    index = 0
                    id_mid = 0
                    for i in pred_boxes.id:
                        # 暂存单次检测中的id最大值
                        if id > id_mid:
                            id_mid = id
                        id = i.to(torch.int)
                        if id > self.max_id:
                            cls = pred_boxes.cls[index]
                            self.count_num[names[cls.item()]] += 1    
                        index += 1           
                    self.max_id = id
                ##################用于计数################
                name = ('' if id is None else f'id:{id} ') + names[c]
                label = (f'{name} {conf:.2f}' if conf else name) if labels else None
                annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True))

代码修改

下面说一下具体修改哪些地方

result.py

# 初始化
 self.count_num = {}
 self.max_id = 0
 for v in self.names.values():
     self.count_num[v] = 0

plot函数

# 新增参数
count_num=None,
max_id=None,

 # 用于计数
 if count_num != None and count_num!={}:
     self.count_num = count_num
 if max_id != None:
     self.max_id = max_id

修改# Plot Detect results 模块：

       # Plot Detect results
       if pred_boxes and show_boxes:
           for d in reversed(pred_boxes):

               c, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
               # 第一帧
               if self.max_id == 0:
                   self.max_id = id
                   for i in pred_boxes.cls:
                       self.count_num[names[i.item()]] += 1 
                ##################用于计数################
                elif id is not None and id > self.max_id:
                    index = 0
                    id_mid = 0
                    for i in pred_boxes.id:
                        # 暂存单次检测中的id最大值
                        if id > id_mid:
                            id_mid = id
                        id = i.to(torch.int)
                        if id > self.max_id:
                           cls = pred_boxes.cls[index]
                           self.count_num[names[cls.item()]] += 1    
                       index += 1           
                   self.max_id = id
               ##################用于计数################
               name = ('' if id is None else f'id:{id} ') + names[c]
               label = (f'{name} {conf:.2f}' if conf else name) if labels else None
               annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True))
       # 将计数结果展示到图片上
       annotator.draw_pic(self.count_num)

其中，以下代码主要是应对第一帧可能存在多目标一次检出的情况。

if self.max_id == 0:
     self.max_id = id
     for i in pred_boxes.cls:
         self.count_num[names[i.item()]] += 1 

下列代码中的for循环遍历每次检出的id列表，并如果该id是首次出现，就通过对应的类别tensor找到对应类别，并且在计数的dict上对应键值加一。 self.max_id = id 用于存储当前遇到的最大id，实测中发现yolov8的id存在问题，并不是一个一个递增的，现在还不知道原因。

  elif id is not None and id > self.max_id:
      index = 0
      id_mid = 0
      for i in pred_boxes.id:
          # 暂存单次检测中的id最大值
          if id > id_mid:
              id_mid = id
          id = i.to(torch.int)
          if id > self.max_id:
              cls = pred_boxes.cls[index]
              self.count_num[names[cls.item()]] += 1  
              print(self.count_num[names[cls.item()]])  
          index += 1           
      self.max_id = id_mid

plotting.py

展示计数结果，放在annotator类里面

    # 用于绘制计数的显示
    def draw_pic(self, count_num):
        class_counts_s = ''
        for k,v in count_num.items():
            class_counts_s += str(k)
            class_counts_s += ':'
            class_counts_s += str(v)
            class_counts_s += '  '
        cv2.putText(self.im, class_counts_s, (30,80), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,0), 3)

predictor.py

这里主要是调用plot函数

 self.count_num = None
 self.max_id = None

write_results函数中进行如下修改：

        if self.args.save or self.args.show:  # Add bbox to image
            plot_args = {
                'line_width': self.args.line_width,
                'boxes': self.args.boxes,
                'conf': self.args.show_conf,
                'labels': self.args.show_labels,
                # 用于计数
                'count_num':self.count_num,
                'max_id':self.max_id
                }
            if not self.args.retina_masks:
                plot_args['im_gpu'] = im[idx]
            self.plotted_img = result.plot(**plot_args)
            # 用于计数
            self.count_num = result.count_num
            self.max_id = result.max_id

部分类别计数

配置文件修改

例如我需要检测coco中的人和椅子，对应的是0和56：

# ultralytics\cfg\default.yaml
classes: [0,56] # (int | list[int], optional) filter results by class, i.e. classes=0, or classes=[0,2,3]

result.py

from ultralytics.cfg import get_cfg
from ultralytics.utils import LOGGER, SimpleClass, ops, DEFAULT_CFG

在前面的代码的基础上，做如下修改：

  args = get_cfg(DEFAULT_CFG)
  classes = args.classes
 # 通过循环取出多个键值对并组装成新字典
 new_dict = {k: self.count_num[k] for i, k in enumerate(self.count_num) if i in classes}
 # 用于计数
 annotator.draw_pic(new_dict)

参考:https://blog.csdn.net/qq_46417453/article/details/130535449