Spaces:
Running
on
T4
Running
on
T4
File size: 1,816 Bytes
06f26d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2022-05-18 09:18:02
import random
import argparse
from pathlib import Path
parser = argparse.ArgumentParser(prog='Face dataset Generation')
parser.add_argument('--face_dir', default='/home/jupyter/data/FFHQ/images512x512', type=str,
metavar='PATH', help="Path to save the face images")
# FFHQ: png, Celeba: png
parser.add_argument('--prefix', default='ffhq', type=str, help="Image format of the HR face images")
parser.add_argument('--num_val', default=500, type=int, help="Ratio for Validation set")
parser.add_argument('--seed', default=1234, type=int, help="Random seed")
parser.add_argument('--im_size', default=512, type=int, help="Random seed")
args = parser.parse_args()
base_dir = Path(__file__).resolve().parents[2] / 'files_txt'
if not base_dir.exists():
base_dir.mkdir()
path_list = sorted([str(x.resolve()) for x in Path(args.face_dir).glob('*.png')])
file_path = base_dir / f"{args.prefix}{args.im_size}.txt"
if file_path.exists():
file_path.unlink()
with open(file_path, mode='w') as ff:
for line in path_list: ff.write(line+'\n')
random.seed(args.seed)
random.shuffle(path_list)
num_train = int(len(path_list) - args.num_val)
file_path_train = base_dir / f"{args.prefix}{args.im_size}_train.txt"
if file_path_train.exists():
file_path_train.unlink()
with open(file_path_train, mode='w') as ff:
for line in path_list[:num_train]: ff.write(line+'\n')
file_path_val = base_dir / f"{args.prefix}{args.im_size}_val.txt"
if file_path_val.exists():
file_path_val.unlink()
with open(file_path_val, mode='w') as ff:
for line in path_list[num_train:]: ff.write(line+'\n')
print('Train / Validation: {:d}/{:d}'.format(num_train, len(path_list)-num_train))
|