Spaces:
Running
on
T4
Running
on
T4
File size: 2,135 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
#!/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')
|