Spaces:
Running
Running
| import math | |
| import os | |
| import numpy as np | |
| from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels | |
| from basicsr.utils import DiffJPEG, USMSharp | |
| from numpy.typing import NDArray | |
| from PIL import Image | |
| from torch import Tensor | |
| from torch.nn import functional as F | |
| from data.degradations import apply_real_esrgan_degradations | |
| from utils.utils_video import img2tensor | |
| blur_kernel_list1 = ['iso', 'aniso', 'generalized_iso', | |
| 'generalized_aniso', 'plateau_iso', 'plateau_aniso'] | |
| blur_kernel_list2 = ['iso', 'aniso', 'generalized_iso', | |
| 'generalized_aniso', 'plateau_iso', 'plateau_aniso'] | |
| blur_kernel_prob1 = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03] | |
| blur_kernel_prob2 = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03] | |
| kernel_size = 21 | |
| blur_sigma1 = [0.05, 0.2] | |
| blur_sigma2 = [0.05, 0.1] | |
| betag_range1 = [0.7, 1.3] | |
| betag_range2 = [0.7, 1.3] | |
| betap_range1 = [0.7, 1.3] | |
| betap_range2 = [0.7, 1.3] | |
| def degrade_imgs(src_folder: str, dst_folder: str, degrade_scale: float, start_size: int) -> None: | |
| src_img_filenames = os.listdir(src_folder) | |
| jpeg_simulator = DiffJPEG() | |
| usm_sharpener = USMSharp() | |
| for src_img_filename in src_img_filenames: | |
| src_img = Image.open(os.path.join(src_folder, src_img_filename)) | |
| src_tensor = img2tensor(np.array(src_img), bgr2rgb=False, | |
| float32=True).unsqueeze(0) / 255.0 | |
| orig_h, orig_w = src_tensor.size()[2:4] | |
| print("SRC TENSOR orig size: ", src_tensor.size()) | |
| if orig_h != start_size or orig_w != start_size: | |
| src_tensor = F.interpolate(src_tensor, size=(start_size, start_size), mode='bicubic') | |
| print("SRC TENSOR new size: ", src_tensor.size()) | |
| blur_kernel1, blur_kernel2, sinc_kernel = _decide_kernels() | |
| (src, src_sharp, degraded_img) = apply_real_esrgan_degradations( | |
| src_tensor, | |
| blur_kernel1=Tensor(blur_kernel1).unsqueeze(0), | |
| blur_kernel2=Tensor(blur_kernel2).unsqueeze(0), | |
| second_blur_prob=0.4, | |
| sinc_kernel=Tensor(sinc_kernel).unsqueeze(0), | |
| resize_prob1=[0.2, 0.7, 0.1], | |
| resize_prob2=[0.3, 0.4, 0.3], | |
| resize_range1=[0.9, 1.1], | |
| resize_range2=[0.9, 1.1], | |
| gray_noise_prob1=0.2, | |
| gray_noise_prob2=0.2, | |
| gaussian_noise_prob1=0.2, | |
| gaussian_noise_prob2=0.2, | |
| noise_range=[0.01, 0.2], | |
| poisson_scale_range=[0.05, 0.45], | |
| jpeg_compression_range1=[85, 100], | |
| jpeg_compression_range2=[85, 100], | |
| jpeg_simulator=jpeg_simulator, | |
| random_crop_gt_size=start_size, | |
| sr_upsample_scale=1, | |
| usm_sharpener=usm_sharpener | |
| ) | |
| # print(src.size()) | |
| # print(src_sharp.size()) | |
| # print(degraded_img.size()) | |
| # print(torch.max(src)) | |
| # print(torch.max(src_sharp)) | |
| # print(torch.max(degraded_img)) | |
| # print(torch.min(src)) | |
| # print(torch.min(src_sharp)) | |
| # print(torch.min(degraded_img)) | |
| # Image.fromarray((src[0] * 255.0).permute(1, 2, 0).cpu().numpy().astype(np.uint8)).save( | |
| # "/home/cll/Desktop/TEST_IMAGE1.png") | |
| # Image.fromarray((src_sharp[0] * 255.0).permute( | |
| # 1, 2, 0).cpu().numpy().astype(np.uint8)).save( | |
| # "/home/cll/Desktop/TEST_IMAGE2.png") | |
| Image.fromarray((degraded_img[0] * 255.0).permute( | |
| 1, 2, 0).cpu().numpy().astype(np.uint8)).save( | |
| os.path.join(dst_folder, src_img_filename)) | |
| print("SAVED %s: " % src_img_filename) | |
| # Image.fromarray((src_tensor[0] * 255.0).permute( | |
| # 1, 2, 0).cpu().numpy().astype(np.uint8)).save( | |
| # os.path.join(dst_folder, src_img_filename)) | |
| # print("SAVED %s: " % src_img_filename) | |
| if __name__ == "__main__": | |
| SRC_FOLDER = "/home/cll/Desktop/sr_test_GT_HQ" | |
| OUTPUT_RESOLUTION_SCALE = 1 | |
| DST_FOLDER = "/home/cll/Desktop/sr_test_degraded_LQ_512" | |
| # DST_FOLDER = "/home/cll/Desktop/sr_test_GT_512" | |
| os.makedirs(DST_FOLDER, exist_ok=True) | |
| degrade_imgs(SRC_FOLDER, DST_FOLDER, OUTPUT_RESOLUTION_SCALE, 512) | |