YOLO v8 在马赛克增强中加入负样本,解决误检问题

0. 前言
虽然v8支持直接训练“无标签的负样本”(v8中叫做background),但默认方式也只是混进正样本中一起练,这会导致新的误检图得不到充分训练
使用本文的“负样本马赛克”功能,可以保证每一轮、每一张图里全都有误检图。能够在较少的训练轮次中,解决漏检问题

使用例:
比如说你发现你模型把摩托车识别成人,那就把那几张摩托车的图放进neg_dir,练个30轮
出来就不会有摩托车误识别的问题了

【但要注意,放进neg_dir里的图 绝 对 不 能 有正样本(例如把摩托车误检成人的图里不能出现真的人)。不然模型会学到错误的负样本,反而会出大问题】

1. 加入新超参
在ultralytics/cfg/default.yaml(默认超参文件)中加入新超参"neg_dir"。
v8的结构是直接读取所有超参数,解析交给后面的函数

neg_dir: '  '  # 负样本文件夹(str),其中只有负样本图片,没有标签(默认为空)

2. 读取新超参
超参在ultralytics/engine/Trainer.py中的BaseTrainer()初始化,
在ultralytics/engine/Trainer.py的108行左右,即BaseTrainer的__init__()方法中,加入下述内容。
目的是读取args里的neg_dir(108行左右改为下述内容),之后超参数会随着trainer的属性hyp一层一层的传下去

 # 负样本文件夹——初始化
 self.neg_dir = self.args.neg_dir

3. 识别超参

在ultralytics/data/augment.py里的v8_transforms()里的Mosaic()。
增加一个输入neg_dir = hyp.neg_dir。

def v8_transforms(dataset, imgsz, hyp, stretch=False):
    """Convert images to a size suitable for YOLOv8 training."""
    pre_transform = Compose([
        Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic, neg_dir=hyp.neg_dir, neg_num = hyp.neg_num),
        CopyPaste(p=hyp.copy_paste),
        RandomPerspective(
            degrees=hyp.degrees,
            translate=hyp.translate,
            scale=hyp.scale,
            shear=hyp.shear,
            perspective=hyp.perspective,
            pre_transform=None if stretch else LetterBox(new_shape=(imgsz, imgsz)),
        )])

再到Mosaic()类中的__init__中,添加识别此参数的代码

    def __init__(self, dataset, imgsz=640, p=1.0, n=4, neg_dir='', neg_num=''):
        """Initializes the object with a dataset, image size, probability, and border."""
        assert 0 <= p <= 1.0, f'The probability should be in range [0, 1], but got {p}.'
        assert n in (4, 9), 'grid must be equal to 4 or 9.'
        super().__init__(dataset=dataset, p=p)
        self.dataset = dataset
        self.imgsz = imgsz
        self.border = (-imgsz // 2, -imgsz // 2)  # width, height
        self.n = n

        # # 更改此处, 读取负样本加入数(默认为2)
        self.neg_num = int(neg_num)
        self.img_neg_files = []  # 负样本路径列表

        # “默认只有负样文件夹才有无标签负样本”
        # 读取负样本文件夹
        # additional feature
        if os.path.isdir(neg_dir):
            # 负样本路径
            self.img_neg_files = [os.path.join(neg_dir, i) for i in os.listdir(neg_dir)]
            logging.info(
                colorstr("Negative dir: ")
                + f"'{neg_dir}', using {len(self.img_neg_files)} picture
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值