Yolov5根据目标类别自定义识别框的颜色_yolov5识别颜色

修改下detect文件下的write results函数

# Write results
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
    label = f'{names[int(cls)]} {conf:.2f}'
    #定义颜色
    color_dict = {'1': [0, 0, 250], '2': [190, 90, 92], '3': [142, 154, 78], '4': [2, 76, 82],
                  '5': [119, 80, 5], '6': [189, 163, 234]}  # Rgb3个值：B G R
    if names[int(cls)] == 'car':        #根据训练的文件.yaml中的name类中的名称修改
         color_single = color_dict['1']
    elif names[int(cls)] == 'person':        
         color_single = color_dict['2']
 
    plot_one_box(xyxy, im0, label=label, color=color_single, line_thickness=3)

参考文档：

(53条消息) 【目标检测】YOLOv5：标签中文显示/自定义颜色_yolov5中文标签_zstar-_的博客-CSDN博客