YOLO-World的微调:提升模型对新类别和不敏感类别的识别性能————保姆级教程

原创

已于 2024-12-26 21:10:03 修改 · 2.5k 阅读

于 2024-12-26 20:35:36 首次发布

最近我在用 YOLO-World 做室内物体检测时，发现一个问题：墙裙这个类别总是检测不出来。不管是用官方提供的预训练权重，还是直接用默认设置，它都没有被识别出来。于是我就在想，能不能在不再把原有数据训练一遍的情况下，增加一个新的类别呢。

考虑到 YOLO-World 结合了 YOLO 检测和 CLIP 跨模态模型，我的思路是利用这两部分的优势进行增强。具体来说，通过收集与新类别（如墙裙）相关的数据集并进行标注，进行单类别训练，让 CLIP 模型通过学习类别的文本描述来识别图像中的新类别特征。这样可以提升 YOLO-World 在实际应用中的新类别识别能力。（如有错误，欢迎批评指正，不胜感激！）

一、数据集准备

我们找好相应类别的数据集后，手动进行标注，这里我使用的是Labelimg进行标注。具体的可以参考这位老哥的帖子：LabelImg（目标检测标注工具）的安装与使用教程-CSDN博客

非常简单，总体的思路就是使用Win+R，运行cmd。输入:

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

然后直接在cmd种输入labelmg启动即可进行对数据的标注，我当时标签是以VOC的形式保存的。

二、数据集的格式转换

YOLO-World官方源码中支持对COCO、Lvis、Object365数据集的训练，在这里我选择把自定义的数据集更新为COCO样式的数据集。首先把所有的图片文件以及打好标签的xml文件放到同一个文件夹下，紧接着用下面的代码来转为将VOC格式数据集转换成COCO格式数据集：

"""
需要修改的地方
1. category_set = ['TBC']此处的类别信息需要修改
2. 代码末尾处的voc标注文件夹
3. 带末尾处对应生成的json文件名
4. 51行处，生成coco标签文件格式后缀名要与自己图片文件类型对应（此处为jpg）
注：训练集与测试集中类别的顺序必须保持一致，因此最好事先确定category的顺序，书写在category_set中
"""

import xml.etree.ElementTree as ET
import os
import json
import collections

coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

# category_set = dict()
image_set = set()
image_id = 1  # train:2018xxx; val:2019xxx; test:2020xxx
category_item_id = 1
annotation_id = 1
'''
  要添加的类别
  '''
category_set = ['wainscot','wall','cabinet']

def addCatItem(name):
    '''
    增加json格式中的categories部分
    '''
    global category_item_id
    category_item = collections.OrderedDict()
    category_item['supercategory'] = 'none'
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_item_id += 1


def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    # image_item = dict()    #按照一定的顺序，这里采用collections.OrderedDict()
    image_item = collections.OrderedDict()
    jpg_name = os.path.splitext(file_name)[0] + '.jpg'
    image_item['file_name'] = jpg_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    image_item['id'] = image_id
    coco['images'].append(image_item)
    image_set.add(jpg_name)
    image_id = image_id + 1
    return image_id


def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    # annotation_item = dict()
    annotation_item = collections.OrderedDict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])
    annotation_item['segmentation'].append(seg)
    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_item['id'] = annotation_id
    annotation_item['ignore'] = 0
    annotation_id += 1
    coco['annotations'].append(annotation_item)


def parseXmlFiles(xml_path):

标签