模拟真实世界的退化代码（python）：RealESRGAN二阶降质代码独立剥离出来使用

一、背景

在超分辨率重建以及修复去噪等领域，如何模拟真实世界的退化对于落地与实际效果尤其重要。在这方面的工作中，RealESRGAN做出了出色的工作。当我们想要借鉴其数据退化思路进行数据集处理与制作时，可以尝试使用其数据处理代码，但是其代码高度耦合，不利于我们使用，于是我将其降质代码剥离出来，可以方便地模拟图像的退化，相对于其他退化方法，这种退化模拟效果较好。

二、退化方法

降质流程的数学表达可简化为如下公式1，其中y表示待降质图像，表示x降质结果，表示D降质算子。其中降质算子包括高斯模糊、噪声、JPEG压缩等。如此，模拟出残片降质的过程，生成大量数据样本对，以训练出用于清晰化的模型,如超分辨率重建等。

（1）

基本降质框架如图所示。通过常见的操作子对图像进行模拟降质，并结合sinc滤波模拟拍摄等因素产生的振铃和伪影等，通过随机扩散操作子模拟纸质老化，最终生成降质结果。

三、代码

代码需要安装pytorch等，我的版本是1.8。

运行如下箭头文件即可：

主要代码：

# coding=gbk
import cv2
import math
import numpy as np
import os
import os.path as osp
import random
import time
import torch
import torchvision
from torch.nn import functional as F
from utils.degradations import circular_lowpass_kernel, random_mixed_kernels,random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from utils import file_client
from utils.logger import get_root_logger
from utils.img_util import imfrombytes,img2tensor
from utils.img_process_util import filter2D,USMSharp
from utils.transforms import augment,paired_random_crop
from utils.DiffJPEG import DiffJPEG
import PIL.Image as Image
class RealESRGANDatadegration():
    def __init__(self,img_paths):
        super(RealESRGANDatadegration, self).__init__()
        self.paths = img_paths
        self.scale = 1 #缩放比例
        #self.gt_size = gt_size #输入图像的尺寸，需要h=w
        gpu_id = None #0
        device = None #
        if gpu_id:
            self.device = torch.device(
                f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
        else:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device

        # the first degradation process
        self.resize_prob = [0.2, 0.7, 0.1]  # up, down, keep
        self.resize_range = [0.15, 1.5]
        self.gaussian_noise_prob = 0.5
        self.noise_range = [1, 30]
        self.poisson_scale_range = [0.05, 3]
        self.gray_noise_prob = 0.4
        self.jpeg_range = [30, 95]

        # the second degradation process
        self.second_blur_prob = 0.8
        self.resize_prob2 = [0.3, 0.4, 0.3]  # up, down, keep
        self.resize_range2 = [0.3, 1.2]
        self.gaussian_noise_prob2 = 0.5
        self.noise_range2 = [1, 25]
        self.poisson_scale_range2 = [0.05, 2.5]
        self.gray_noise_prob2 = 0.4
        self.jpeg_range2 = [30, 95]

        # blur settings for the first degradation
        self.blur_kernel_size = 21
        self.kernel_list = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
        self.kernel_prob =[0.45, 0.25, 0.12, 0.03, 0.12, 0.03]  # a list for each kernel probability
        self.sinc_prob = 0.1
        self.blur_sigma = [0.2, 3]
        self.betag_range = [0.5, 4] # betag used in generalized Gaussian blur kernels
        self.betap_range = [1, 2] # betap used in plateau blur kernels

        ## blur settings for the second degradation
        self.blur_kernel_size2 =21
        self.kernel_list2 = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
        self.kernel_prob2 = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
        self.sinc_prob2 = 0.1
        self.blur_sigma2 = [0.2, 1.5]
        self.betag_range2 = [0.5, 4]
        self.betap_range2 = [1, 2]

        # a final sinc filter
        self.final_sinc_prob = 0.8
        ############################
        self.use_hflip = False
        self.use_rot = False
        ################
        self.kernel_range = [2 * v + 1 for v in range(3, 11)] ## kernel size ranges from 7 to 21
        # TODO: kernel range is now hard-coded, should be in the configure file
        self.pulse_tensor = torch.zeros(21, 21).float()  # convolving with pulse tensor brings no blurry effect
        self.pulse_tensor[10, 10] = 1

        self.file_client = None
        self.usm_sharpener = USMSharp().to(self.device)  # do usm sharpening
        self.jpeger = DiffJPEG(differentiable=False).to(self.device) # # simulate JPEG compression artifacts
    @torch.no_grad()
    def kernel(self):
        if self.file_client is None:
            self.file_client = file_client.FileClient()

        # -------------------------------- Load gt images -------------------------------- #
        # Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
        gt_path = self.paths
        # avoid errors caused by high latency in reading files
        retry = 3
        while retry > 0:
            try:
                img_bytes = self.file_client.get(gt_path, 'gt')
            except (IOError, OSError) as e:
                logger = get_root_logger()
                logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
                # change another file to read
                index = random.randint(0, self.__len__())
                gt_path = self.paths[index]
                time.sleep(1)  # sleep 1s for occasional server congestion
            else:
                break
            finally:
                retry -= 1

        img_gt = imfrombytes(img_bytes, float32=True)
        #
        # # -------------------- Do augmentation for training: flip, rotation -------------------- #
        # img_gt = augment(img_gt, self.use_hflip, self.use_rot)
        #
        # # crop or pad to 400
        # # TODO: 400 is hard-coded. You may change it accordingly
        # h, w = img_gt.shape[0:2]
        # crop_pad_size = self.gt_size
        # # pad
        # if h < crop_pad_size or w < crop_pad_size:
        #     pad_h = max(0, crop_pad_size - h)
        #     pad_w = max(0, crop_pad_size - w)
        #     img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
        # # crop
        # if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
        #     h, w = img_gt.shape[0:2]
        #     # randomly choose top and left coordinates
        #     top = random.randint(0, h - crop_pad_size)
        #     left = random.randint(0, w - crop_pad_size)
        #     img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]

        # ------------------------ Generate kernels (used in the first degradation) ------------------------ #
        kernel_size = random.choice(self.kernel_range)
        if np.random.uniform() < self.sinc_prob:
            # 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(
                self.kernel_list,
                self.kernel_prob,
                kernel_size,
                self.blur_sigma,
                self.blur_sigma, [-math.pi, math.pi],
                self.betag_range,
                self.betap_range,
                noise_range=None)
        # pad kernel
        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(self.kernel_range)
        if np.random.uniform() < self.sinc_prob2:
            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(
                self.kernel_list2,
                self.kernel_prob2,
                kernel_size,
                self.blur_sigma2,
                self.blur_sigma2, [-math.pi, math.pi],
                self.betag_range2,
                self.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)))

        # ------------------------------------- the final sinc kernel ------------------------------------- #
        if np.random.uniform() < self.final_sinc_prob:
            kernel_size = random.choice(self.kernel_range)
            omega_c = np.random.uniform(np.pi / 3, np.pi)
            sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
            sinc_kernel = torch.FloatTensor(sinc_kernel)
        else:
            sinc_kernel = self.pulse_tensor

        # BGR to RGB, HWC to CHW, numpy to tensor
        img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
        # img_gt = img_gt.unsqueeze(4)
        # print(img_gt.shape)
        kernel = torch.FloatTensor(kernel)
        kernel2 = torch.FloatTensor(kernel2)

        return_d = {'gt': img_gt, 'kernel1': kernel, 'kernel2': kernel2, 'sinc_kernel': sinc_kernel, 'gt_path': gt_path}
        return return_d

    @torch.no_grad()
    def synthesis(self):
        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
        """
        # training data synthesis
        data = self.kernel()
        gt = data['gt'].to(self.device)
        self.gt = gt.unsqueeze(0)
        self.gt_usm = self.usm_sharpener(self.gt)
        self.kernel1 = data['kernel1'].to(self.device)
        self.kernel2 = data['kernel2'].to(self.device)
        self.sinc_kernel = data['sinc_kernel'].to(self.device)

        ori_h, ori_w = self.gt.size()[2:4]

        # ----------------------- The first degradation process ----------------------- #
        # blur
        out = filter2D(self.gt_usm, self.kernel1)
        # random resize
        updown_type = random.choices(['up', 'down', 'keep'], self.resize_prob)[0]
        if updown_type == 'up':
            scale = np.random.uniform(1, self.resize_range[1])
        elif updown_type == 'down':
            scale = np.random.uniform(self.resize_range[0], 1)
        else:
            scale = 1
        mode = random.choice(['area', 'bilinear', 'bicubic'])
        out = F.interpolate(out, scale_factor=scale, mode=mode)
        # add noise
        gray_noise_prob = self.gray_noise_prob
        if np.random.uniform() < self.gaussian_noise_prob:
            out = random_add_gaussian_noise_pt(
                out, sigma_range=self.noise_range, clip=True, rounds=False, gray_prob=gray_noise_prob)
        else:
            out = random_add_poisson_noise_pt(
                out,
                scale_range=self.poisson_scale_range,
                gray_prob=gray_noise_prob,
                clip=True,
                rounds=False)
        # JPEG compression
        jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.jpeg_range)
        out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
        out = self.jpeger(out, quality=jpeg_p)

        # ----------------------- The second degradation process ----------------------- #
        # blur
        if np.random.uniform() < self.second_blur_prob:
            out = filter2D(out, self.kernel2)
        # random resize
        updown_type = random.choices(['up', 'down', 'keep'], self.resize_prob2)[0]
        if updown_type == 'up':
            scale = np.random.uniform(1, self.resize_range2[1])
        elif updown_type == 'down':
            scale = np.random.uniform(self.resize_range2[0], 1)
        else:
            scale = 1
        mode = random.choice(['area', 'bilinear', 'bicubic'])
        out = F.interpolate(
            out, size=(int(ori_h / self.scale * scale), int(ori_w / self.scale * scale)),
            mode=mode)
        # add noise
        gray_noise_prob = self.gray_noise_prob2
        if np.random.uniform() < self.gaussian_noise_prob2:
            out = random_add_gaussian_noise_pt(
                out, sigma_range=self.noise_range2, clip=True, rounds=False, gray_prob=gray_noise_prob)
        else:
            out = random_add_poisson_noise_pt(
                out,
                scale_range=self.poisson_scale_range2,
                gray_prob=gray_noise_prob,
                clip=True,
                rounds=False)

        # JPEG compression + the final sinc filter
        # We also need to resize images to desired sizes. We group [resize back + sinc filter] together
        # as one operation.
        # We consider two orders:
        #   1. [resize back + sinc filter] + JPEG compression
        #   2. JPEG compression + [resize back + sinc filter]
        # Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
        if np.random.uniform() < 0.5:
            # resize back + the final sinc filter
            mode = random.choice(['area', 'bilinear', 'bicubic'])
            out = F.interpolate(out, size=(ori_h // self.scale, ori_w // self.scale), mode=mode)
            out = filter2D(out, self.sinc_kernel)
            # JPEG compression
            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.jpeg_range2)
            out = torch.clamp(out, 0, 1)
            out = self.jpeger(out, quality=jpeg_p)
        else:
            # JPEG compression
            jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.jpeg_range2)
            out = torch.clamp(out, 0, 1)
            out = self.jpeger(out, quality=jpeg_p)
            # resize back + the final sinc filter
            mode = random.choice(['area', 'bilinear', 'bicubic'])
            out = F.interpolate(out, size=(ori_h // self.scale, ori_w // self.scale), mode=mode)
            out = filter2D(out, self.sinc_kernel)

        # clamp and round
        self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.

        # # random crop
        # gt_size = self.gt_size
        # (self.gt, self.gt_usm), self.lq = paired_random_crop([self.gt, self.gt_usm], self.lq, gt_size,
        #                                                      self.scale)

        # training pair pool
        #self._dequeue_and_enqueue()
        # sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
        self.gt_usm = self.usm_sharpener(self.gt)
        lq = self.lq.contiguous()
        lq = lq.squeeze(0)
        # lq = lq.cpu().numpy()
        # #lq = 255 * (1.0 - lq)
        # lq = Image.fromarray(lq.astype(np.uint8), mode='RGB')
        return lq
if __name__ == '__main__':
    #img = torch.rand(3,256,256)
    t = RealESRGANDatadegration('1.png')
    result = t.synthesis()
    torchvision.utils.save_image(result, '2.png')

四、结果：

输入：

realesrgan二阶降质效果输出1：

realesrgan二阶降质效果输出2：

结果：

第一行为原图，第二行为双三次插值，第三行为realesrgan二阶降质效果

五、code

https://download.csdn.net/download/Crystal_remember/87396238

链接：百度网盘 请输入提取码 提取码：Luck