使用yolov5只检测一种类别

使用yolov5只检测一种类别

参考文章：(102条消息) yolov5只检测单一类别或者特定的类_雪可问春风的博客-CSDN博客_yolov5单分类

使用作者在GitHub中给出的训练好的权重文件。

权重文件地址下载地址

本文仅用了yolov5s.pt做试验

最初的yolov5检测结果：

代码修改

在detect.py中 （我的代码里是在228行）不知道的可以ctrl + F 搜索一下

原代码

 if save_img or save_crop or view_img:  # Add bbox to image
     c = int(cls)  # integer class
     label = None if hide_labels else (names[c] if hide_conf     else f'{names[c]} {conf:.2f}')
     annotator.box_label(xyxy, label, color=colors(c, True))
     if save_crop:
         save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

修改为：

 if save_img or save_crop or view_img:  # Add bbox to image
     c = int(cls)  # integer class
     # FIXME:
     if names[int(cls)] == "name_sign":
         label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
         annotator.box_label(xyxy, label, color=colors(c, True),counta = counta)
         if save_crop:
               save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

再运行结果：