Spaces:
Running
on
T4
Running
on
T4
File size: 6,341 Bytes
561c629 |
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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
# -*- coding: utf-8 -*-
import argparse
import cv2
import numpy as np
import os, shutil
import sys
from multiprocessing import Pool
from os import path as osp
from tqdm import tqdm
import random
from collections import namedtuple
# Import files from the local folder
root_path = os.path.abspath('.')
sys.path.append(root_path)
from degradation.ESR.usm_sharp import USMSharp
class worker:
def __init__(self, start_index=1):
# The index you want to start with
self.output_index = start_index
def process(self, path, opt, usm_sharper):
''' crop the image here (also do usm here)
Args:
path (str): path of the image
opt (dict): all setting in a dictionary
usm_sharper (class): usm sharpener
Returns:
cropped_num (int): how many cropped images you have for this path
'''
crop_size = opt['crop_size'] # usually 400
# read image
img = cv2.imread(path)
height, width = img.shape[0:2]
res_store = []
crop_num = (height//crop_size)*(width//crop_size)
random_num = opt['crop_num_per_img']
# Use shift offset to make image more cover origional image size
shift_offset_h, shift_offset_w = 0, 0
if random_num == -1:
# We should select all sub-frames order by order (not randomly select here)
choices = [i for i in range(crop_num)]
shift_offset_h = 0 #random.randint(0, height - crop_size * (height//crop_size))
shift_offset_w = 0 #random.randint(0, width - crop_size * (width//crop_size))
else:
# Divide imgs by crop_size x crop_size and choose opt['crop_num_per_img'] num of them to avoid overlap
num = min(random_num, crop_num)
choices = random.sample(range(crop_num), num)
for choice in choices:
row_num = (width//crop_size)
x, y = crop_size * (choice // row_num), crop_size * (choice % row_num)
# add offset
res_store.append((x, y))
# Sharp the image before selection
if opt['usm_save_folder'] != None:
sharpened_img = usm_sharper(img)
for (h, w) in res_store:
cropped_img = img[h+shift_offset_h : h+crop_size+shift_offset_h, w+shift_offset_w : w+crop_size+shift_offset_w, ...]
cropped_img = np.ascontiguousarray(cropped_img)
cv2.imwrite(osp.join(opt['save_folder'], f'img_{self.output_index:06d}.png'), cropped_img, [cv2.IMWRITE_PNG_COMPRESSION, 0]) # Save in lossless mode
# store the sharpened cropped image
if opt['usm_save_folder'] != None:
cropped_sharpened_img = sharpened_img[h+shift_offset_h : h+crop_size+shift_offset_h, w+shift_offset_w : w+crop_size+shift_offset_w, ...]
cropped_sharpened_img = np.ascontiguousarray(cropped_sharpened_img)
cv2.imwrite(osp.join(opt['usm_save_folder'], f'img_{self.output_index:06d}.png'), cropped_sharpened_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
self.output_index += 1
cropped_num = len(res_store)
return cropped_num
def extract_subimages(opt):
# Input
input_folders = opt['input_folders']
# Make folders
save_folder = opt['save_folder']
usm_save_folder = opt['usm_save_folder']
if osp.exists(save_folder):
print(f'Folder {save_folder} already exists. Program will delete this folder!')
shutil.rmtree(save_folder)
os.makedirs(save_folder)
if usm_save_folder != None:
if osp.exists(usm_save_folder):
print(f'Folder {usm_save_folder} already exists. Program will delete this folder!')
shutil.rmtree(usm_save_folder)
print("Use usm sharp")
os.makedirs(usm_save_folder)
# USM
usm_sharper = USMSharp(type="cv2")
# Iterate all datasets' folders
start_index = 1
for input_folder in input_folders:
print(input_folder, start_index)
# Scan all images
img_list = []
for file in sorted(os.listdir(input_folder)):
if file.split(".")[-1] in ["png", "jpg"]:
img_list.append(osp.join(input_folder, file))
# Iterate can crop
obj = worker(start_index=start_index) # The start_index determines where you will start your naming your image (usually start from 0)
for path in img_list:
if random.random() < opt['select_rate']:
cropped_num = obj.process(path, opt, usm_sharper)
start_index += cropped_num
print(start_index, path)
else:
print("SKIP")
print('All processes done.')
def main(args):
opt = {}
input_folders = []
if type(args.input_folder) == str:
input_folders.append(args.input_folder)
else:
for input_folder in args.input_folder:
input_folders.append(input_folder)
print("input folders have ", input_folders)
opt['input_folders'] = input_folders
opt['save_folder'] = args.save_folder
opt['usm_save_folder'] = args.output_usm
opt['crop_size'] = args.crop_size
opt['crop_num_per_img'] = args.crop_num_per_img
opt['select_rate'] = args.select_rate
# Extract subimages
extract_subimages(opt)
if __name__ == '__main__':
random.seed(777) # We setup a random seed such that all program get the same cropped images
parser = argparse.ArgumentParser()
# Try to split image after default
parser.add_argument('-i', '--input_folder', nargs='+', type=str, default='datasets/all_Anime_hq_frames_resize', help='Input folder') # TODO: support multiple image input
parser.add_argument('-o', '--save_folder', type=str, default='datasets/train_hr', help='Output folder')
parser.add_argument('--output_usm', type=str, help='usm sharpened hr folder')
parser.add_argument('--crop_size', type=int, default=360, help='Crop size')
parser.add_argument('--select_rate', type=float, default=1, help='(0-1): Proportion to keep; 1 means to keep them all')
parser.add_argument('--crop_num_per_img', type=int, default=-1, help='Crop size (int); -1 means use all possible sub-frames')
args = parser.parse_args()
main(args) |