YOLOv8 Bytetrack部分类别计数
基于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