CUB_200_2011 数据集预处理批量 crop 裁剪 + split 划分 python 实现

数据集

CUB-200-2011
200类鸟，共11788张图片
官方介绍页面
下载：CUB-200-2011 dataset by classes folder

任务描述

根据可解释深度学习论文 《This Looks Like That: Deep Learning for Interpretable Image Recognition》 的代码中Readme.txt文件要求预处理数据集（裁剪图片 + 划分数据集）

必要的文件结构

图中蓝色的 CUB_200_2011 文件夹与 prepare_data.py 同级；
其中 CUB_200_2011 文件夹由数据集压缩包 CUB_200_2011.tgz 解压而成；
prepare_data.py 内容见下文。

代码实现

import os
import pandas as pd
from PIL import Image
from shutil import copyfile

def makedir(path):
    '''
    if path does not exist in the file system, create it
    '''
    if not os.path.exists(path):
        os.makedirs(path)

# set paths
rootpath = 'CUB_200_2011/CUB_200_2011/'
imgspath = rootpath + 'images/'
trainpath = 'datasets/cub200_cropped/train_cropped/'
testpath = 'datasets/cub200_cropped/test_cropped/'

# read img names, bounding_boxes
names = pd.read_table(rootpath + 'images.txt', delimiter=' ', names=['id', 'name'])
names = names.to_numpy()
boxs = pd.read_table(rootpath + 'bounding_boxes.txt', delimiter=' ',
                     names=['id', 'x', 'y', 'width', 'height'])
boxs = boxs.to_numpy()

# crop imgs
for i in range(11788):
    im = Image.open(imgspath + names[i][1])
    im = im.crop((boxs[i][1], boxs[i][2], boxs[i][1] + boxs[i][3], boxs[i][2] + boxs[i][4]))
    im.save(imgspath + names[i][1], quality=95)
    print('{} imgs cropped and saved.'.format(i + 1))
print('All Done.')

# mkdir for cropped imgs
folders = pd.read_table(rootpath + 'classes.txt', delimiter=' ', names=['id', 'folder'])
folders = folders.to_numpy()
for i in range(200):
    makedir(trainpath + folders[i][1])
    makedir(testpath + folders[i][1])

# split imgs
labels = pd.read_table(rootpath + 'train_test_split.txt', delimiter=' ', names=['id', 'label'])
labels = labels.to_numpy()
for i in range(11788):
    if(labels[i][1] == 1):
        copyfile(imgspath + names[i][1], trainpath + names[i][1])
    else:
        copyfile(imgspath + names[i][1], testpath + names[i][1])
    print('{} imgs splited.'.format(i + 1))
print('All Done.')

预处理结果

生成与蓝色的 CUB_200_2011 文件夹同级的 datasets 文件夹

一些值得注意的地方

1、PIL.Image.crop()方法的裁剪参数设置参考《python PIL库的crop函数–图片裁剪操作》、数据集 Readme.txt 中如下部分。

=========================
BOUNDING BOXES:
=========================

Each image contains a single bounding box label.  Bounding box labels are contained in the file bounding_boxes.txt, with each line corresponding to one image:

<image_id> <x> <y> <width> <height>

where <image_id> corresponds to the ID in images.txt, and <x>, <y>, <width>, and <height> are all measured in pixels

2、PIL.Image.save()方法中quality参数设置参考《Python Pillow (PIL) Image.save 保存为jpg图片压缩问题》
默认quality=75会压缩图片；
在quality=95时为原图文件体积；
quality=100时比原图文件大小更大。