pix2pix pytorch 学习笔记（①）数据集处理 （详细） |
您所在的位置：网站首页 › 数据集处理网站 › pix2pix pytorch 学习笔记（①）数据集处理 （详细） |
pix2pix pytorch 学习笔记（①）数据集处理 （详细）
|
最近在学习深度学习的相关知识，阅读了经典的pix2pix论文： 论文原址：[1611.07004] Image-to-Image Translation with Conditional Adversarial Networks (arxiv.org) 下面是我学习论文和代码部分的总结，因为没有很多代码基础所以进度缓慢。这次是代码部分； 1.从原材料数据开始： 提供了facade库，里面有606张图片；
需要对图片文件进行分类和处理。才能获得训练所用的数据集。 由于文件中提供的是名字相同jpg与png格式的图片，所以我对源代码进行了修改，以能够进行图像的拼接。 我的文件例分类，暂时不分为test, train , val ;先使用 combine_A_and_B.py 文件进行处理整个文件。
文件夹A为 数据集里的jpg文件，即真实图形，文件夹B为png文件，即彩色轮廓图。文件夹dataset是预置的存放生成训练的图像文件夹。另外要注意在A，B的文件夹下还要各建立一个相同名的子文件夹，图像文件存放其中。我这里取名为base。 存放完毕后，学习并修改代码。 import os import numpy as np import cv2 import argparse import multiprocessing from multiprocessing import Pool if __name__ == '__main__': multiprocessing.freeze_support() # 确保多进程功能在冻结的可执行文件中正常工作 def image_write(path_A, path_B, path_AB): im_A = cv2.imread(path_A, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR im_B = cv2.imread(path_B, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR im_AB = np.concatenate([im_A, im_B], 1) cv2.imwrite(path_AB, im_AB) # 定义命令行指令 parser = argparse.ArgumentParser('create image pairs') parser.add_argument('--fold_A', dest='fold_A', help='input directory for image A', type=str, default='../dataset/50kshoes_edges') # 输入A文件夹文件 parser.add_argument('--fold_B', dest='fold_B', help='input directory for image B', type=str, default='../dataset/50kshoes_jpg') parser.add_argument('--fold_AB', dest='fold_AB', help='output directory', type=str, default='../dataset/test_AB') parser.add_argument('--num_imgs', dest='num_imgs', help='number of images', type=int, default=1000000) parser.add_argument('--use_AB', dest='use_AB', help='if true: (0001_A, 0001_B) to (0001_AB)', action='store_true') parser.add_argument('--no_multiprocessing', dest='no_multiprocessing', help='If used, chooses single CPU execution instead of parallel execution', action='store_true', default=False) args = parser.parse_args() for arg in vars(args): print('[%s] = ' % arg, getattr(args, arg)) splits = os.listdir(args.fold_A) if not args.no_multiprocessing: pool = Pool() # 创建一个进程池对象“pool”，用于并行处理数据；未使用 --no_multiprocessing 参数，则会执行进程池的创建。提高数据处理的效率 for sp in splits: # 对于”splits“下的每一个子文件 img_fold_A = os.path.join(args.fold_A, sp) img_fold_B = os.path.join(args.fold_B, sp) # 构建源图像文件夹路径 img_list = os.listdir(img_fold_A) # 获取图像文件夹中的文件列表 if args.use_AB: img_list = [img_path for img_path in img_list if 'cmb_' in img_path] # 使用此操作的话将筛选文件名含有‘_A.’的文件 num_imgs = min(args.num_imgs, len(img_list)) # 计算实际使用的图像数量，取到最小值，控制合并的图像数量 print('split = %s, use %d/%d images' % (sp, num_imgs, len(img_list))) img_fold_AB = os.path.join(args.fold_AB, sp) # 输出图像的文件夹路径 if not os.path.isdir(img_fold_AB): os.makedirs(img_fold_AB) # 不存在的话就创建输出的此文件夹 print('split = %s, number of images = %d' % (sp, num_imgs)) # 合并循环； for n in range(num_imgs): img_name = os.path.splitext(img_list[n])[0] # 这里将jpg格式的名字单独取出来 name_A = img_list[n] path_A = os.path.join(img_fold_A, name_A) if args.use_AB: name_B = name_A.replace('_A.', '_B.') # 检查是否使用了arg.use_AB else: name_B = f"{img_name}.png" path_B = os.path.join(img_fold_B, name_B) # 构建B图像文件的完整路径 if os.path.isfile(path_A) and os.path.isfile(path_B): name_AB = name_A if args.use_AB: name_AB = name_AB.replace('_A.', '.') # remove _A path_AB = os.path.join(img_fold_AB, name_AB) if not args.no_multiprocessing: pool.apply_async(image_write, args=(path_A, path_B, path_AB)) else: im_A = cv2.imread(path_A, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR im_B = cv2.imread(path_B, 1) # python2: cv2.CV_LOAD_IMAGE_COLOR; python3: cv2.IMREAD_COLOR im_AB = np.concatenate([im_A, im_B], 1) cv2.imwrite(path_AB, im_AB) if not args.no_multiprocessing: pool.close() pool.join()代码如上，参考junyanz/pytorch-CycleGAN-and-pix2pix: Image-to-Image Translation in PyTorch (github.com)的combine_A_and_B.py 进行了修改。仅供参考，欢迎提出改进意见。 执行时可以在终端中使用 python datasets/combine_A_and_B.py --fold_A /path/to/data/A --fold_B /path/to/data/B --fold_AB /path/to/datafold_A, fold_B, fold_AB 请自行根据自己实际情况修改。最后可以得到的拼接图像如下
这样就得到了用于训练的数据集了。 欢迎交流。 |
【本文地址】
公司简介
联系我们
今日新闻 |
点击排行 |
|
推荐新闻 |
图片新闻 |
|
专题文章 |