【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理

协助学长做车辆数据集(下载)扩充，从车辆的一个角度生成车辆的另外一个角度，使用GAN网络。目前的GAN网络的输入为512*256的图像，下图为网络输入样例，左边是生成图像，右边是原始图像。将数据集用GAN网络训练后，网络可以根据右边的图像生成左边的图像。现在做的是根据原始数据集，生成如下图所示的pair-wise数据集。
这里写图片描述
现有的车辆原始数据集如下图所示(样例，并非全部)：
数据集1（目录1_5）：

数据集2(目录2_5)：

我们需要生成如下pair-wise成对图像（样例），作为GAN网络的训练样本：

为生成512*256的pair-wise图像，我们需要进行如下两步：1. 图像尺寸规格化； 2. 不同方向成对图像拼接（保持车ID不变）

一、图像尺寸规格化
观察图像可知，每一张图像的尺寸大小都不同，而输入图像尺寸是确定的，为512*256（两张图像拼接），亦即每一张图像为256*256。因此首先对图像作尺寸规格化，代码如下：

from PIL import Image
import os.path
import glob
def convertjpg(jpgfile,outdir,width=256,height=256):
    img=Image.open(jpgfile)
    try:
        new_img=img.resize((width,height),Image.BILINEAR)   
        new_img.save(os.path.join(outdir,os.path.basename(jpgfile)))
    except Exception as e:
        print(e)

for i in range(8):
    dir_s="D:\\大创项目\数据集\\Im2Im\\dataset(copy)\\dataset\\"+str(i+1)+"_5"+"\\*.jpg"
    dir_d="D:\\大创项目\\数据集\\Im2Im\\dataset(new)\\"+str(i+1)+"_5"
    for jpgfile in glob.glob(dir_s):
        convertjpg(jpgfile,dir_d)

尺寸规格化后的数据集如下图所示，每一张图像的尺寸都是256*256：
这里写图片描述

二、不同方向成对图像拼接（保持车ID不变）
生成pair-wise数据集，并按照train：val：test=5：2：3的比例分配生成训练集、验证集和测试集。
共生成八对数据集，分别为前后左右四个方向的相邻两个方向的成对数据集。
另外一点，对于数据集1（生成后在右边的图像，或称为原始图像）中存在而数据集2（生成后在左边的图像，或称为目标图像）中不存在的车辆图片，将其复制两边成pair-wise图像放入测试集中。也起到了数据增强的效果。（测试只利用右边的原始图像生成左边的目标图像）。代码如下：

from PIL import Image
import os.path
import glob
import re

def Merge(jpgfile_left,jpgfile_right,outdir,count):
    arr = [jpgfile_left,jpgfile_right]
    toImage = Image.new('RGB',(512,256))
    for i in range(2):
        fromImge = Image.open(arr[i])
        loc = ((i % 2) * 256 , (int(i/2) * 256))
        #print(loc)
        toImage.paste(fromImge, loc)
    jpgfile_right=jpgfile_right[:-4]+"_merge"+str(count)+".jpg"
    toImage.save(os.path.join(outdir,os.path.basename(jpgfile_right)))

ar=[1,1,3,3,5,5,7,7];
br=[2,8,4,2,6,4,8,6];

for i in range(8):
    ##获取图像的路径##
    dir_right="D:\\大创项目\\数据集\\Im2Im\\dataset(new)\\"+str(ar[i])+"_5\\*.jpg"
    dir_left="D:\\大创项目\\数据集\\Im2Im\\dataset(new)\\"+str(br[i])+"_5\\*.jpg"
    ##获取图像的路径##

    ##建立文件夹##
    root_dir="D:\\大创项目\\数据集\\Im2Im\\merge("+str(ar[i])+"_"+str(br[i])+")"
    os.mkdir(root_dir)
    train_dir=root_dir+"\\train"
    test_dir=root_dir+"\\test"
    val_dir=root_dir+"\\val"
    os.mkdir(train_dir)
    os.mkdir(test_dir)
    os.mkdir(val_dir)
    ##建立文件夹##

    ##计算配对图像个数##
    ks=0;
    for jpgfile_right in glob.glob(dir_right):
        #print(jpgfile_right)
        result_right=re.findall('.*?(\d+)\_c',jpgfile_right)[0];
        #print(result_right)
        for jpgfile_left in glob.glob(dir_left):
            result_left=re.findall('.*?(\d+)\_c',jpgfile_left)[0];
            if result_left==result_right:
                ks+=1;
    print(ks)
    ##计算配对图像个数##

    ##配对并生成拼接图像##
    k=0;
    count2=0;
    for jpgfile_right in glob.glob(dir_right):
        #print(jpgfile_right)
        result_right=re.findall('.*?(\d+)\_c',jpgfile_right)[0];
        #print(result_right)
        count=0;
        for jpgfile_left in glob.glob(dir_left):
            result_left=re.findall('.*?(\d+)\_c',jpgfile_left)[0];
            if result_left==result_right:
                count+=1
                k+=1;
                if k<=ks*0.5:
                    Merge(jpgfile_left,jpgfile_right,train_dir,count)
                elif k<=ks*0.7:
                    Merge(jpgfile_left,jpgfile_right,val_dir,count)
                else :
                    Merge(jpgfile_left,jpgfile_right,test_dir,count)
                if count==5:
                    break
                continue
            if result_left>result_right:##不存在配对图像（单向），则图片复制拼接
                count2+=1;
                Merge(jpgfile_right,jpgfile_right,test_dir,count2)
                if count2==5:
                    count2=0
                break
    ##配对并生成拼接图像##

生成了8个目录，目录格式为merge(n_m)【n,m分别为数据集中的子数据集idx】。每一个目录下含有train\val\test三个数据集。
某一数据集的样例如下图所示：
这里写图片描述

【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理

猜你喜欢