stylegan3-标签制作脚本

最新推荐文章于 2023-12-31 01:09:26 发布

ye1256

最新推荐文章于 2023-12-31 01:09:26 发布

阅读量217

点赞数

文章标签： python 深度学习 pytorch

本文链接：https://blog.csdn.net/ye1256/article/details/129603545

版权

该脚本用于生成StyleGAN2-ADA-PyTorch的多类别JSON标签文件。它从输入文件夹读取数据，为每个文件创建标签，并将结果保存到指定的输出文件夹中的dataset.json。程序使用argparse模块处理命令行参数，os模块遍历文件系统，以及json模块进行序列化。

摘要由CSDN通过智能技术生成


import argparse
import os
import random
import json

def parse_args():
	desc = "Tool to create multiclass json labels file for stylegan2-ada-pytorch" 
	parser = argparse.ArgumentParser(description=desc)

	parser.add_argument('--verbose', action='store_true',
		help='Print progress to console.')

	parser.add_argument('--input_folder', type=str,
		default=r'C:\Users\yewenjing\Desktop\GAN_input1',
		help='Directory path to the inputs folder. (default: %(default)s)')

	parser.add_argument('--output_folder', type=str,
		default=r'C:\Users\yewenjing\Desktop\GAN_input1',
		help='Directory path to the outputs folder. (default: %(default)s)')

	args = parser.parse_args()
	return args

def main():
	global args
	args = parse_args()
	base_dir = ''

	remakePath = args.output_folder
	if not os.path.exists(remakePath):
		os.makedirs(remakePath)

	data_dict = {}
	data_dict['labels'] = []
	label_counter = 0

	with open(os.path.join(remakePath, 'dataset.json'), 'w') as outfile:

		for root, subdirs, files in os.walk(args.input_folder):
			if len(subdirs) > 0:
				base_dir = root	
				continue

			current_subdir = os.path.split(root)[1]

			for filename in files:
				file_path = os.path.join(current_subdir, filename).replace('\\','/')
				if(args.verbose): print('\t- file %s (full path: %s)' % (filename, file_path))
				data_dict['labels'].append([file_path, label_counter])
				
			label_counter += 1

		json.dump(data_dict, outfile)


if __name__ == "__main__":
	main()