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

协助学长做车辆数据集(下载)扩充,从车辆的一个角度生成车辆的另外一个角度,使用GAN网络。目前的GAN网络的输入为512*256的图像,下图为网络输入样例,左边是生成图像,右边是原始图像。将数据集用GAN网络训练后,网络可以根据右边的图像生成左边的图像。现在做的是根据原始数据集生成如下图所示的pair-wise数据集
【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理_第1张图片
现有的车辆原始数据集如下图所示(样例,并非全部):
数据集1(目录1_5)
【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理_第2张图片
数据集2(目录2_5)
【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理_第3张图片
我们需要生成如下pair-wise成对图像(样例),作为GAN网络的训练样本:
【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理_第4张图片

为生成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:
【数据增强】基于GAN网络生成车辆不同角度图像之训练数据集处理_第5张图片

二、不同方向成对图像拼接(保持车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网络生成车辆不同角度图像之训练数据集处理_第6张图片

你可能感兴趣的:(-----项目记录-----)