Spaces:
Running
on
T4
Running
on
T4
File size: 1,805 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 |
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2022-05-18 07:58:01
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parents[3]))
import argparse
import multiprocessing
import albumentations as Aug
from utils import util_image
parser = argparse.ArgumentParser(prog='SISR dataset Generation')
parser.add_argument('--face_dir', default='/home/jupyter/data/FFHQ/images1024x1024', type=str,
metavar='PATH', help="Path to save the HR face images")
parser.add_argument('--save_dir', default='/home/jupyter/data/FFHQ/', type=str,
metavar='PATH', help="Path to save the resized face images")
# FFHQ: png
parser.add_argument('--ext', default='png', type=str, help="Image format of the HR face images")
parser.add_argument('--pch_size', default=512, type=int, metavar='PATH', help="Cropped patch size")
args = parser.parse_args()
# check the floder to save the cropped patches
pch_size = args.pch_size
pch_dir = Path(args.face_dir).parent / f"images{pch_size}x{pch_size}"
if not pch_dir.exists(): pch_dir.mkdir(parents=False)
transform = Aug.Compose([Aug.SmallestMaxSize(max_size=pch_size),])
# HR face path
path_hr_list = [x for x in Path(args.face_dir).glob('*.'+args.ext)]
def process(im_path):
im = util_image.imread(im_path, chn='rgb', dtype='uint8')
pch = transform(image=im)['image']
pch_path = pch_dir / (im_path.stem + '.png')
util_image.imwrite(pch, pch_path, chn='rgb')
num_workers = multiprocessing.cpu_count()
pool = multiprocessing.Pool(num_workers)
pool.imap(func=process, iterable=path_hr_list, chunksize=16)
pool.close()
pool.join()
num_pch = len([x for x in pch_dir.glob('*.png')])
print('Totally process {:d} images'.format(num_pch))
|