Python/torch/深度学习——VOC数据集图像去白边

Python/torch/深度学习——VOC数据集图像去白边

---

---

前言

VOC-2012的语义分割数据集存在白边，需要去除，本文提供了一种去白边的函数

---

一、VOC语义分割数据集下载

VOC数据集总共有21类（包含背景，num_classes=21就对了）

0: 背景（Background）
1: 飞机（Aeroplane）
2: 自行车（Bicycle）
3: 鸟（Bird）
4: 船（Boat）
5: 瓶子（Bottle）
6: 巴士（Bus）
7: 汽车（Car）
8: 猫（Cat）
9: 椅子（Chair）
10: 牛（Cow）
11: 餐桌（Dining Table）
12: 狗（Dog）
13: 马（Horse）
14: 摩托车（Motorbike）
15: 人（Person）
16: 盆栽植物（Potted Plant）
17: 羊（Sheep）
18: 沙发（Sofa）
19: 火车（Train）
20: 电视（TV / Monitor）

VOC-2012年版数据集可以通过 torchvision 直接下载，或是使用下载链接直接下载，后期要额外处理，因此这里可以不使用 transforms
VOC-2012年数据集下载链接

import torchvision


VOC_data = torchvision.datasets.VOCSegmentation(# 保存地址文件夹
                                                      root=r"F:\torch_cuda_117\DataFromNet", 
                                                       
                                                      # 2007或2012版，这里是2012版
                                                      year='2012',  
                                                      
                                                      # 配置是训练集还是数据集
                                                      image_set='train',  
                                                      
                                                      # 添加transforms，可选
                                                      # transform=VOC_transforms,  
                                                      # target_transform=VOC_transforms,
                                                      
                                                      download=True)

---

二、VOC语义分割数据集解压文件夹简介

下载后的压缩包为 VOCtrainval_11-May-2012.tar ，解压后得到 VOCdevkit 进入到 VOCdevkit\VOC2012 可以看到5个文件夹：

• Annotations 为xml文件
• ImageSets 为文本文件（目前不管，有人说是要根据里面的说明选择文件，这里并没有这么做）
• JPEGImages 为输入原图
• SegmentationClass 为语义分割分类标注（在这里使用的文件夹）
• SegmentationObject 为实例分割标注

---

三、数据处理

一般情况下会使用 DataLoader(VOC_train_data, batch_size=4, shuffle=True) 这个命令直接加载，但是在实际使用中出现了各种状况，导致训练出现问题，同时发现标注有白边需要去除，因此对要使用的数据进行转存和处理

1、数据转存

这里根据 SegmentationClass 中有的文件进行复制并转存

import os
import shutil

def DataReSave(base_folder, source_folder, target_folder) :
    """
    基于base_folder,将source_folder中同名文件存到target_folder
    :param base_folder: 文件名来源
    :param source_folder: 要筛选的文件所在文件夹
    :param target_folder: 目标文件夹
    """
    names = os.listdir(base_folder)
    for name in names:
        source_path = os.path.join(source_folder, name)  # .replace("png", "jpg")
        target_path = os.path.join(target_folder, name)  # .replace("png", "jpg")
        if os.path.exists(source_path):
            shutil.copyfile(source_path, target_path)  # 将图片另存到目标文件夹
            print(f"图片 {name} 已保存")
        else:
            print(f"图片 {name} 不存在于源文件夹中")


DataReSave(base_folder, source_folder, target_folder)

2.去白边并保存

去白边后看不到labels的显示图像，但是能用代码查看，如果要肉眼观察，需要进行后处理，在这不管
代码如下：

def VOC_data_noWhite(source_path, target_path):
    """
    图像去白边
    :param source_path: 原文件目录
    :param target_path: 目标文件目录
    :return: 去白边后的文件，与原文件同名
    """
    imgs = os.listdir(source_path)
    for name in imgs:
        img = Image.open(os.path.join(source_path, name))
        # img_resize = img.resize(size)
        img_np = np.array(img)
        # print(set(img_np.reshape(-1)))
        img_np_noWhite = np.where(img_np == 255, 0, img_np)
        # print(set(img_np_noWhite.reshape(-1)))  # 查看是否已经去掉白边
        img_PIL_noWhite = Image.fromarray(img_np_noWhite)
        # img_PIL_noWhite.show()
        img_PIL_noWhite.save(os.path.join(target_path, name))