Spaces:
Running
on
T4
Running
on
T4
File size: 4,042 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 |
# -*- coding: utf-8 -*-
'''
From ESRGAN
'''
import os, sys
import cv2
import numpy as np
import torch
from torch.nn import functional as F
from scipy import special
import random
import math
from torchvision.utils import make_grid
from degradation.ESR.degradations_functionality import *
root_path = os.path.abspath('.')
sys.path.append(root_path)
def np2tensor(np_frame):
return torch.from_numpy(np.transpose(np_frame, (2, 0, 1))).unsqueeze(0).cuda().float()/255
def tensor2np(tensor):
# tensor should be batch size1 and cannot be grayscale input
return (np.transpose(tensor.detach().squeeze(0).cpu().numpy(), (1, 2, 0))) * 255
def mass_tensor2np(tensor):
''' The input tensor is massive tensor
'''
return (np.transpose(tensor.detach().squeeze(0).cpu().numpy(), (0, 2, 3, 1))) * 255
def save_img(tensor, save_name):
np_img = tensor2np(tensor)[:,:,16]
# np_img = np.expand_dims(np_img, axis=2)
cv2.imwrite(save_name, np_img)
def filter2D(img, kernel):
"""PyTorch version of cv2.filter2D
Args:
img (Tensor): (b, c, h, w)
kernel (Tensor): (b, k, k)
"""
k = kernel.size(-1)
b, c, h, w = img.size()
if k % 2 == 1:
img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
else:
raise ValueError('Wrong kernel size')
ph, pw = img.size()[-2:]
if kernel.size(0) == 1:
# apply the same kernel to all batch images
img = img.view(b * c, 1, ph, pw)
kernel = kernel.view(1, 1, k, k)
return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
else:
img = img.view(1, b * c, ph, pw)
kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
def generate_kernels(opt):
kernel_range = [2 * v + 1 for v in range(opt["kernel_range"][0], opt["kernel_range"][1])]
# ------------------------ Generate kernels (used in the first degradation) ------------------------ #
kernel_size = random.choice(kernel_range)
if np.random.uniform() < opt['sinc_prob']:
# 里面加一层sinc filter，但是10%的概率
# this sinc filter setting is for kernels ranging from [7, 21]
if kernel_size < 13:
omega_c = np.random.uniform(np.pi / 3, np.pi)
else:
omega_c = np.random.uniform(np.pi / 5, np.pi)
kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
else:
kernel = random_mixed_kernels(
opt['kernel_list'],
opt['kernel_prob'],
kernel_size,
opt['blur_sigma'],
opt['blur_sigma'], [-math.pi, math.pi],
opt['betag_range'],
opt['betap_range'],
noise_range=None)
# pad kernel: -在v2我是直接省略了padding
pad_size = (21 - kernel_size) // 2
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
# ------------------------ Generate kernels (used in the second degradation) ------------------------ #
kernel_size = random.choice(kernel_range)
if np.random.uniform() < opt['sinc_prob2']:
# 里面加一层sinc filter，但是10%的概率
if kernel_size < 13:
omega_c = np.random.uniform(np.pi / 3, np.pi)
else:
omega_c = np.random.uniform(np.pi / 5, np.pi)
kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
else:
kernel2 = random_mixed_kernels(
opt['kernel_list2'],
opt['kernel_prob2'],
kernel_size,
opt['blur_sigma2'],
opt['blur_sigma2'], [-math.pi, math.pi],
opt['betag_range2'],
opt['betap_range2'],
noise_range=None)
# pad kernel
pad_size = (21 - kernel_size) // 2
kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
kernel = torch.FloatTensor(kernel)
kernel2 = torch.FloatTensor(kernel2)
return (kernel, kernel2)
|