CV图像分类数据集处理

数据集保存格式1

• train
  • class1
    • 0.jpg
    • 1.jpg
  • class2
    • 0.jpg
    • 1.jpg
• test
  • 0.jpg
  • 1.jpg
    图片按类别分别保存在不同文件夹下。

划分数据集

import codecs
import os
import random
import shutil
from PIL import Image

train_ratio = 4 / 5
all_file_dir = '/home/aistudio/data/all_data/train'
class_list = [c for c in os.listdir(all_file_dir) if os.path.isdir(os.path.join(all_file_dir, c)) and not c.endswith('Set') and not c.startswith('.')]
class_list.sort()
print(class_list)
train_image_dir = os.path.join(all_file_dir, "trainImageSet")
if not os.path.exists(train_image_dir):
    os.makedirs(train_image_dir)
    
eval_image_dir = os.path.join(all_file_dir, "evalImageSet")
if not os.path.exists(eval_image_dir):
    os.makedirs(eval_image_dir)

train_file = codecs.open(os.path.join(all_file_dir, "train.txt"), 'w')
eval_file = codecs.open(os.path.join(all_file_dir, "eval.txt"), 'w')

with codecs.open(os.path.join(all_file_dir, "label_list.txt"), "w") as label_list:
    label_id = 0
    for class_dir in class_list:
        label_list.write("{0}\t{1}\n".format(label_id, class_dir))
        image_path_pre = os.path.join(all_file_dir, class_dir)
        for file in os.listdir(image_path_pre):
            try:
                img = Image.open(os.path.join(image_path_pre, file))
                if random.uniform(0, 1) <= train_ratio:
                    shutil.copyfile(os.path.join(image_path_pre, file), os.path.join(train_image_dir, file))
                    src = os.path.join(train_image_dir, file) #原先的图片名字
                    dst = os.path.join(train_image_dir, str(class_dir) + file) #根据自己的需要重新命名,可以把'E_' + img改成你想要的名字
                    os.rename(src, dst) #重命名,覆盖原先的名字

                    train_file.write("{0} {1}\n".format(os.path.join("trainImageSet", str(class_dir) + file), label_id))
                else:
                    shutil.copyfile(os.path.join(image_path_pre, file), os.path.join(eval_image_dir, file))
                    src = os.path.join(eval_image_dir, file) #原先的图片名字
                    dst = os.path.join(eval_image_dir, str(class_dir) + file) #根据自己的需要重新命名,可以把'E_' + img改成你想要的名字
                    os.rename(src, dst) #重命名,覆盖原先的名字
                    
                    eval_file.write("{0} {1}\n".format(os.path.join("evalImageSet", str(class_dir) + file), label_id))
            except Exception as e:
                pass
                # 存在一些文件打不开，此处需要稍作清洗
        label_id += 1
            
train_file.close()
eval_file.close()

在train文件夹下生成trainImageSet、evalImageSet、train.txt、eval.txt、label_list.txt四个文件。

数据保存格式2

• train
  • trainImageSet
    • 0.jpg
    • 1.jpg
  • evalImageSet
    • 0.jpg
    • 1.jpg
  • train.txt
  • eval.txt
• test
  • 0.jpg
  • 1.jpg

划分数据集

data_dir = '/home/aistudio/data/new_train'
test_data_dir = '/home/aistudio/data/all_data/test'
train_list = '/home/aistudio/data/new_train/train.txt'
trainval_list = []
with open(train_list) as f:
    lines = f.readlines()
    for line in lines:
        line = line.strip().replace('\t',' ').split(' ')
        trainval_list.append([line[0], int(line[1])])
print(trainval_list[-10:])
#查看训练集中总样本数
print(len(trainval_list))

train_data, val_data = [], []
#划分训练集和验证集比例
train_ratio = 0.8
for i in range(len(trainval_list)):
    if random.uniform(0, 1) <= train_ratio:
        train_data.extend([trainval_list[i]])
    else:
        val_data.extend([trainval_list[i]])
print(train_data[0],train_data[0])
print(len(train_data), len(val_data))

数据增强

PROB = 1 

def random_brightness(img):
    prob = np.random.uniform(0, 1)
    if prob < PROB:
        brightness_delta = 0.125
        delta = np.random.uniform(-brightness_delta, brightness_delta) + 1
        img = ImageEnhance.Brightness(img).enhance(delta)
    return img

def random_contrast(img):
    prob = np.random.uniform(0, 1)
    if prob < PROB:
        contrast_delta = 0.5
        delta = np.random.uniform(-contrast_delta, contrast_delta) + 1
        img = ImageEnhance.Contrast(img).enhance(delta)
    return img

def random_saturation(img):
    prob = np.random.uniform(0, 1)
    if prob < PROB:
        saturation_delta = 0.5
        delta = np.random.uniform(-saturation_delta, saturation_delta) + 1
        img = ImageEnhance.Color(img).enhance(delta)
    return img

def random_hue(img):
    prob = np.random.uniform(0, 1)
    if prob < PROB:
        hue_delta = 18
        delta = np.random.uniform(-hue_delta, hue_delta)
        img_hsv = np.array(img.convert('HSV'))
        img_hsv[:, :, 0] = img_hsv[:, :, 0] + delta         #在HSV格式下，修改图片的亮度V
        img = Image.fromarray(img_hsv, mode='HSV').convert('RGB')
    return img

def expand_image(img, percent):
    w = img.size[0]
    h = img.size[1]
    new_w = int(w * percent)
    new_h = int(h * percent)
    img=img.resize((new_w,new_h),Image.ANTIALIAS)
    return img

def resize_img(img):
    ret = img.resize((227, 227), Image.ANTIALIAS)  #PIL Image.resize抗锯齿的选项
    return ret

def keep_ration_expand(img):#保持长宽比放大，使得最短边为227
    w=img.size[0]
    h=img.size[1]
    percent = max(227.0/w, 227.0/h)
    img = expand_image(img, percent) 
    return img

def center_crop_img(img):
    w=img.size[0]
    h=img.size[1]
    if(w<227 or h<227):
        img = keep_ration_expand(img)
    w=img.size[0]
    h=img.size[1]
    xmin = (w-227)//2
    ymin = (h-227)//2
    xmax = xmin+227
    ymax = ymin+227
    img = img.crop((xmin, ymin, xmax, ymax))
    return img

def flip_left_right(img):
    prob = np.random.uniform(0, 1)
    if prob < 0.5:
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
    return img

def distort_image(img):                 # 随机修改图片的亮度、对比度、饱和度以及HSV格式下的亮度
    prob = np.random.uniform(0, 1)
    # Apply different distort order
    if prob > 0.5:
        img = random_brightness(img)
        img = random_contrast(img)
        img = random_saturation(img)
        img = random_hue(img)
        #img = center_crop_img(img)
        img = flip_left_right(img)
    else:
        img = random_brightness(img)
        img = random_saturation(img)
        img = random_hue(img)
        img = random_contrast(img)
        #img = center_crop_img(img)
        img = flip_left_right(img)
    return img

def img_norm(img_file,outdir,j,label,enhance_train_list):#img_file 是图片绝对路径
    try:
        img = Image.open(img_file)
    except OSError:
        return j
    w = img.size[0]
    h = img.size[1]
    if(w==h):#如果原图是正方形的，就将其缩放至227，然后进行对比度、左右翻转等操作，生成20张图片
        img = img.resize((227, 227), Image.ANTIALIAS)
        #print('w:',w),
        #print('h:',h)
        keep = img
        for i in range(20):#生成20个随机处理的图
            img = keep
            img = distort_image(img)
            #img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
    elif(w>h):#如果原图的宽大于高，则先保持比例缩放，使得最短边等于227，然后分别取图片的左、中、右3个部分进行操作，共生成30张图片
        img = keep_ration_expand(img)
        w = img.size[0]
        h = img.size[1]
        #print('w:',w),
        #print('h:',h)
        keep = img
        for i in range(10):#left
            img = keep
            xmin= 0
            ymin= 0
            xmax= h
            ymax= h
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
        for i in range(10):#center
            img = keep
            xmin= (w-h)//2
            ymin= 0
            xmax= (w+h)//2
            ymax= h
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
        for i in range(10):#right
            img = keep
            xmin= w-h
            ymin= 0
            xmax= w
            ymax= h
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
    else:#w<h #如果原图的宽小于高，则先保持比例缩放，使得最短边等于227，然后分别取图片的上、中、下3个部分进行操作，共生成30张图片
        img = keep_ration_expand(img)
        w = img.size[0]
        h = img.size[1]
        #print('w:',w),
        #print('h:',h)
        keep = img
        for i in range(10):#up
            img = keep
            xmin= 0
            ymin= 0
            xmax= w
            ymax= w
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
        for i in range(10):#center
            img = keep
            xmin= 0
            ymin= (h-w)//2
            xmax= w
            ymax= (h+w)//2
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
        for i in range(10):#down
            img = keep
            xmin= 0
            ymin= h-w
            xmax= w
            ymax= h
            img = img.crop((xmin,ymin,xmax,ymax))
            img = distort_image(img)
            img = img.resize((227, 227), Image.ANTIALIAS)
            out_name = img_file.split('/')[-1].replace(".jpg","")+'_out_'+str(j)+"_"+str(label)+".jpg"
            #mydir = input_img.replace(input_img.split('/')[-1], "")
            save_name = os.path.join(outdir,out_name)
            img.save(save_name)
            with open(enhance_train_list,'a') as f: 
                f.write(save_name+"\t"+label+"\n")
            j+=1
    return j#全局图片编号
    
start = time.time() 
print('start')

outdir  = "/home/aistudio/data/new_train"
old_list = "/home/aistudio/data/all_data/train/train.txt"

if os.path.exists(outdir): #删除已有文件夹
    shutil.rmtree(outdir)
if not os.path.exists(outdir):#重新创建输出文件夹
    os.makedirs(outdir)
j=0

enhance_train_list = "/home/aistudio/work/enhance_train_list.txt"
with open(enhance_train_list,'w') as f: #清空已有文件
    pass

#将图片路径与标签信息加入词典，为后面遍历图片文件夹做准备
dict={} 
with open(old_list,'r') as f: 
    for line in f.readlines():
        img_path, label = line.strip().split(' ')
        img_path = "/home/aistudio/data/all_data/train/"+ img_path
        im = Image.open(img_path)
        if im.mode!='RGB' or im.format!='JPEG':
            im=im.convert('RGB')
            out = img_path.split('.')[0]+"_out.jpg"
            im.save(out) 
            dict[out]=label
        else: 
            dict[img_path]=label

old_train_dir = "/home/aistudio/data/all_data/train/trainImageSet"

#按照文件名的方式写入list，相当于进行了混洗
for _, _, files in os.walk(old_train_dir):
    for img in files:
        absolute_path = os.path.join(old_train_dir,img)
        im = Image.open(absolute_path) 
        if im.mode=='RGB' and im.format=='JPEG':
            label=dict[absolute_path] 
            j=img_norm(absolute_path, outdir, j, label,enhance_train_list)
        else: 
            out = absolute_path.split('.')[0]+"_out.jpg"
            label = dict[out]
            im=im.convert('RGB')
            im.save(out) #im.save(out,format='JPEG')
            j=img_norm(out, outdir, j, label,enhance_train_list)
            

''' 这段代码会使得训练集按照标签顺序写入list,没有进行混洗，训练效果不好
with open(old_list,'r') as f: 
    for line in f.readlines():
        img_path, label = line.strip().split('\t')
        img_path = "/home/aistudio/data/data10954/"+ img_path
        im = Image.open(img_path)
        if im.mode!='RGB' or im.format!='JPEG':
            im=im.convert('RGB')
            out = img_path.split('.')[0]+"_out.jpg"
            im.save(out)#im.save(out,format='JPEG')
            j=img_norm(out, outdir, j, label,enhance_train_list)
        else:
            j=img_norm(img_path, outdir, j, label,enhance_train_list)
'''
print('enhance and gen list, cost time: %.4f s' % (time.time()-start))