mmcls多标签分类实战(一):制作多标签数据

Multi-Label:多标签分类指当图片或者文本存在多个属性时,多标签分类能够对一张图片中的多个属性进行分类识别。下面将用mmclassification框架完成一次多标签分类实战。

Multi-class与Multi-label的区别
首先,我们需要构建数据集,需要注意的是,多标签分类的数据集与多分类数据不同。Multi-class的label中只存在一个‘1’,而Multi-label中对应多个属性,label可以有多个‘1’。Multi-class构建数据集比较简单,只需要将不同类别的图片归纳到不同文件下即可。Multi-Label数据集需要将label写成txt,并自定义dataset去读取。
mmcls多标签分类实战(一):制作多标签数据_第1张图片

Multi-label构建数据集
构建Multi-label数据集也比较简单。首先,需要将图片针对属性分类,将图片分别归纳到不同属性名称的文件夹下(由于一张图片可能存在多个属性,所以同一张图片可以存放在不同属性文件夹下)。我们用字典来实现图片与属性的对应。遍历所有文件夹下的图片,将图片名作为字典的key,文件名作为字典的value。遍历完后,就可以获得图片属性的label,保存在txt中。

import os

file_handle = open('val.txt', mode='w')
img_dic = {}
label2idx = {'LH':0, 'PS':1, 'SL':2, 'WH':3, 'ZC':4}

for home, dirs, files in os.walk('/home/mmclassification-master/data/multi_val/'):
    for filename in files:
        if filename not in img_dic:
            img_dic[filename] = [label2idx[home.split('/')[-1]]]
        else:
            img_dic[filename].append(label2idx[home.split('/')[-1]])
     
for k, v in img_dic.items():
    file_handle.writelines(os.path.join("/home/mmclassification-master/data/multi_val",\
                             k) + ' ' + str(v)[1:-1].replace(" ", "") +'\n')
file_handle.close()

构建的txt文件大概就如下图所示,每一行包含图片名称与其对应的属性标签。下面,我们将自定义dataset来读取数据。
mmcls多标签分类实战(一):制作多标签数据_第2张图片

你可能感兴趣的:(pytorch,计算机视觉,深度学习,分类)