| |
| |
| |
|
|
| 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') |
|
|
|
|
