Spaces:
Sleeping
Sleeping
#!/usr/bin/env python | |
# -*- coding:utf-8 -*- | |
# Power by Zongsheng Yue 2022-07-16 12:42:42 | |
import sys | |
from pathlib import Path | |
sys.path.append(str(Path(__file__).resolve().parents[3])) | |
import os | |
import math | |
import torch | |
import argparse | |
from einops import rearrange | |
from datapipe.datasets import DatasetBicubic | |
from utils import util_image | |
from utils import util_common | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
"--files_txt", | |
type=str, | |
default='./datapipe/files_txt/celeba512_val.txt', | |
help="File names") | |
parser.add_argument( | |
"--sf", | |
type=int, | |
default=8, | |
help="Number of trainging iamges", | |
) | |
parser.add_argument( | |
"--bs", | |
type=int, | |
default=8, | |
help="Batch size", | |
) | |
parser.add_argument( | |
"--save_dir", | |
type=str, | |
default='', | |
help="Folder to save the fake iamges", | |
) | |
parser.add_argument( | |
"--num_images", | |
type=int, | |
default=100, | |
help="Number of iamges", | |
) | |
args = parser.parse_args() | |
save_dir = Path(args.save_dir) | |
if not save_dir.stem.endswith(f'x{args.sf}'): | |
save_dir = save_dir.parent / f"{save_dir.stem}_x{args.sf}" | |
util_common.mkdir(save_dir, delete=True) | |
dataset = DatasetBicubic( | |
files_txt=args.files_txt, | |
up_back=True, | |
need_gt_path=True, | |
sf=args.sf, | |
length=args.num_images, | |
) | |
dataloader = torch.utils.data.DataLoader( | |
dataset, | |
batch_size=args.bs, | |
drop_last=False, | |
num_workers=4, | |
pin_memory=False, | |
) | |
for ii, data_batch in enumerate(dataloader): | |
im_lq_batch = data_batch['lq'] | |
im_path_batch = data_batch['gt_path'] | |
print(f"Processing: {ii+1}/{math.ceil(len(dataset) / args.bs)}...") | |
for jj in range(im_lq_batch.shape[0]): | |
im_lq = rearrange( | |
im_lq_batch[jj].clamp(0.0, 1.0).numpy(), | |
'c h w -> h w c', | |
) | |
im_name = Path(im_path_batch[jj]).name | |
im_path = save_dir / im_name | |
util_image.imwrite(im_lq, im_path, chn='rgb', dtype_in='float32') | |