Spaces:
Running
Running
File size: 4,205 Bytes
2514fb4 |
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 |
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)
|