3.YUAN_deblurring_gaussion_Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from threading import Thread  # needed since the denoiser is running in parallel\n",
    "import queue\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim\n",
    "from models.skip import skip  # our network\n",
    "from utils.utils import *  # auxiliary functions\n",
    "from utils.blur_utils import *  # blur functions\n",
    "from utils.data import Data  # class that holds img, psnr, time\n",
    "from skimage.restoration import denoise_nl_means\n",
    "from dncnn_models.network_dncnn import DnCNN as net_dncnn       # dncnn net\n",
    "from dncnn_models.network_ffdnet import FFDNet as net_ffdnet    # FFDNet\n",
    "from utils import utils_logger   # taken from zhangkai\n",
    "from utils import utils_image as util\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cudnn.benchmark = True\n",
    "dtype = torch.cuda.FloatTensor\n",
    "\n",
    "# 参数设置部分：\n",
    "NOISE_SIGMA=25\n",
    "# 两种选择的模糊核：\n",
    "BLUR_TYPE = 'gauss_blur'  # 'gauss_blur' or 'uniform_blur' that the two only options\n",
    "\n",
    "GRAY_SCALE = False  # if gray scale is False means we have rgb image, the psnr will be compared on Y. ch.\n",
    "                    # if gray scale is True it will turn rgb to gray scale\n",
    "USE_FOURIER = False\n",
    "X_LABELS = ['Iterations']*3\n",
    "Y_LABELS = ['PSNR between x and net (db)', 'PSNR with original image (db)', 'loss']\n",
    "ORIGINAL    = 'Clean'\n",
    "CORRUPTED   = 'Blurred'\n",
    "DIP_NLM     = 'DRED (NLM)'\n",
    "DIP_TEMP    = 'DIP_TEMP'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_network_and_input(img_shape, input_depth=32, pad='reflection',\n",
    "                          upsample_mode='bilinear', use_interpolate=True, align_corners=False,\n",
    "                          act_fun='LeakyReLU', skip_n33d=128, skip_n33u=128, skip_n11=4,\n",
    "                          num_scales=5, downsample_mode='stride', INPUT='noise'):  # 'meshgrid'\n",
    "    \"\"\" Getting the relevant network and network input (based on the image shape and input depth)\n",
    "        We are using the same default params as in DIP article\n",
    "        img_shape - the image shape (ch, x, y)\n",
    "    \"\"\"\n",
    "    n_channels = img_shape[0]\n",
    "    net = skip(input_depth, n_channels,\n",
    "               num_channels_down=[skip_n33d] * num_scales if isinstance(skip_n33d, int) else skip_n33d,\n",
    "               num_channels_up=[skip_n33u] * num_scales if isinstance(skip_n33u, int) else skip_n33u,\n",
    "               num_channels_skip=[skip_n11] * num_scales if isinstance(skip_n11, int) else skip_n11,\n",
    "               upsample_mode=upsample_mode, use_interpolate=use_interpolate, align_corners=align_corners,\n",
    "               downsample_mode=downsample_mode, need_sigmoid=True, need_bias=True, pad=pad, act_fun=act_fun).type(dtype)\n",
    "    net_input = get_noise(input_depth, INPUT, img_shape[1:]).type(dtype).detach()\n",
    "    return net, net_input\n",
    "\n",
    "def train_via_admm(net, net_input, denoiser_function, H, y,             # H is the kernel, y is the blurred image\n",
    "                   clean_img=None, plot_array={}, algorithm_name=\"\",             # clean_img for psnr to be shown\n",
    "                   gamma=.9, step_size=1000, save_path=\"\",         # scheduler parameters and path to save params\n",
    "                   admm_iter=3000, LR=0.004,                                          # admm_iter is step_2_iter\n",
    "                   sigma_f=NOISE_SIGMA, update_iter=10, method='fixed_point',  # method: 'fixed_point' or 'grad' or 'mixed'\n",
    "                   LR_x=None, noise_factor=0.01):  # LR_x needed only if method!=fixed_point\n",
    "    # get optimizer and loss function:\n",
    "    mse = torch.nn.MSELoss().type(dtype)  # using MSE loss\n",
    "    # additional noise added to the input:\n",
    "    net_input_saved = net_input.detach().clone()\n",
    "    noise = net_input.detach().clone()\n",
    "    # x update method:\n",
    "    if method == 'fixed_point':\n",
    "        swap_iter = admm_iter + 1\n",
    "        LR_x = None\n",
    "    elif method == 'grad':\n",
    "        swap_iter = -1\n",
    "    elif method == 'mixed':\n",
    "        swap_iter = admm_iter // 2\n",
    "    else:\n",
    "        assert False, \"method can be 'fixed_point' or 'grad' or 'mixed' only \"\n",
    "\n",
    "    # run RED via ADMM, initialize:\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=LR)  # using ADAM opt\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, gamma=gamma, step_size=step_size)\n",
    "    y_torch = np_to_torch(y).type(dtype)\n",
    "    x = y.copy()\n",
    "    f_x, u = x.copy(), np.zeros_like(x)\n",
    "    avg = np.rint(y)\n",
    "    \n",
    "    psnr_net_list=[]\n",
    "    psnr_hr=[]\n",
    "    image_list=[]\n",
    "    psnr_avg_list=[]\n",
    "\n",
    "    # ADMM:\n",
    "    for i in range(1, 1 + admm_iter):\n",
    "\n",
    "        if NOISE_SIGMA == 50:  #optimal\n",
    "            beta=0.02               #optimal\n",
    "            mu=0.04                 #optimal\n",
    "            increase_reg= 0.0001    #optimal\n",
    "            mu = mu + increase_reg\n",
    "            beta = beta + increase_reg\n",
    "        elif NOISE_SIGMA == 25:\n",
    "            beta=0.02\n",
    "            mu=0.02\n",
    "            increase_reg= 0\n",
    "            mu = mu - increase_reg\n",
    "            beta = beta - increase_reg\n",
    "            # 考虑：因为随着迭代噪声变小，减少sigma_f的影响\n",
    "            sigma_f = sigma_f -0.00001\n",
    "\n",
    "        elif NOISE_SIGMA == 15: \n",
    "            beta=0.02\n",
    "            mu=0.04\n",
    "            increase_reg= 0.000001\n",
    "            mu = mu + increase_reg\n",
    "            beta = beta + increase_reg\n",
    "            # 考虑：因为随着迭代噪声变小，减少sigma_f的影响\n",
    "            sigma_f = sigma_f -0.0000001\n",
    "        else:\n",
    "            print('we have not tested other noise level')    \n",
    "    \n",
    "        # step 1, update network:\n",
    "        optimizer.zero_grad()\n",
    "        net_input = net_input_saved + (noise.normal_() * noise_factor)\n",
    "        out = net(net_input)\n",
    "        out_np = torch_to_np(out)\n",
    "        # loss:\n",
    "        loss_y = mse(H(out), y_torch)\n",
    "        loss_x = mse(out, np_to_torch(x - u).type(dtype))\n",
    "        total_loss = loss_y + mu * loss_x\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        # step 2, update x using a denoiser and result from step 1\n",
    "        f_x = denoiser_function(x.copy(), sigma_f)\n",
    "\n",
    "        # 使用深度先验的话需要去掉一维：\n",
    "        f_x=np.squeeze(f_x)\n",
    "\n",
    "        if i < swap_iter:\n",
    "            x = 1 / (beta + mu) * (beta * f_x + mu * (out_np + u))\n",
    "        else:\n",
    "            x = x - LR_x * (beta * (x - f_x) + mu * (x - out_np - u))\n",
    "        np.clip(x, 0, 1, out=x)  # making sure that image is in bounds\n",
    "        # step 3, update u\n",
    "        u = u + out_np - x\n",
    "        # Averaging:\n",
    "        avg = avg * .99 + out_np * .01\n",
    "        # show psnrs:\n",
    "        if clean_img is not None:\n",
    "            psnr_net = compare_PSNR(clean_img, out_np, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "            psnr_x_u = compare_PSNR(clean_img, x - u, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "            psnr_avg = compare_PSNR(clean_img, avg, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "            print('\\r', algorithm_name, '%04d/%04d Loss %f' % (i, admm_iter, total_loss.item()),\n",
    "                  'psnrs: net: %.2f x-u: %.2f avg: %.2f' % (psnr_net, psnr_x_u,psnr_avg), end='')\n",
    "            \n",
    "            # 记录结果：\n",
    "            psnr_hr.append(psnr_x_u)\n",
    "            psnr_net_list.append(psnr_net)\n",
    "            image_list.append(x-u)\n",
    "            psnr_avg_list.append(psnr_avg)\n",
    "\n",
    "            if plot_array:  # plot graphs only if plotting images, since it time-consuming\n",
    "                psnr_x_u = compare_PSNR(clean_img, x - u, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)  \n",
    "                psnr_net = compare_PSNR(clean_img, out_np, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "        else:\n",
    "            print('\\r', algorithm_name, 'iteration %04d/%04d Loss %f' % (i, admm_iter, total_loss.item()), end='')\n",
    "\n",
    "    return x - u, psnr_hr,psnr_net_list,image_list,psnr_avg_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 处理rgb图像的FFDNet先验：\n",
    "\n",
    "n_channels_color=3  #处理单通道图像\n",
    "# 实例化：\n",
    "model_ffdnet_color=net_ffdnet(in_nc=n_channels_color,out_nc=n_channels_color,nc=96,nb=12,act_mode='R')\n",
    "#预训练的FFDNet模型参数：(使用绝对路径！)\n",
    "model_path_ffdnet_color='/home/yuanweimin/PHD_3/2019_ICCVW_DeepRED/model_zoo/ffdnet_color.pth'\n",
    "\n",
    "#加载训练参数：\n",
    "model_ffdnet_color.load_state_dict(torch.load(model_path_ffdnet_color),strict=True)\n",
    "model_ffdnet_color.eval()\n",
    "#通常在实际代码中，在预测阶段，也会加上torch.no_grad()来关闭梯度的计算\n",
    "for k, v in model_ffdnet_color.named_parameters():\n",
    "    v.requires_grad = False\n",
    "\n",
    "def FFDNet_color_yuan(noisy_np_img,sigma):\n",
    "    noisy_torch_img=np_to_torch(noisy_np_img)\n",
    "\n",
    "    n_channels=noisy_torch_img.shape[1]\n",
    "    denoised_img=[]\n",
    "    sigma_map=torch.full((1,1,1,1),sigma/255.).type_as(noisy_torch_img) # size:[1,1,1,1]\n",
    "    #print(noisy_torch_img.shape)  # size:[1,3,256,256]\n",
    "\n",
    "    \n",
    "    # print(noisy_torch_img[:,c,:,:].shape)                 # tensor, size:[1,256,256]\n",
    "    #temp=torch.unsqueeze(noisy_torch_img[:,c,:,:],dim=0)  # tensor, size:[1,1,256,256] (增加一维)\n",
    "    \n",
    "    denoise_torch_fast=model_ffdnet_color(noisy_torch_img,sigma_map)                # denoise_torch_fast:[1,1,256,256]\n",
    "    #print(denoise_torch_fast.shape)\n",
    "\n",
    "    denoise_np_fast=torch_to_np(denoise_torch_fast)       # denoise_np_fast:[1,256,256]\n",
    "    #print(denoise_np_fast.shape)\n",
    "    denoised_img+=[denoise_np_fast]   \n",
    "    return np.array(denoised_img, dtype=np.float32)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_imgs_deblurring(fname, blur_type, noise_sigma, plot=False):\n",
    "    \"\"\"  Loads an image, and add gaussian blur\n",
    "    Args: \n",
    "         fname: path to the image\n",
    "         blur_type: 'uniform' or 'gauss'\n",
    "         noise_sigma: noise added after blur\n",
    "         covert2gray: should we convert to gray scale image?\n",
    "         plot: will plot the images\n",
    "    Out:\n",
    "         dictionary of images and dictionary of psnrs\n",
    "    \"\"\"\n",
    "    img_pil, img_np = load_and_crop_image(fname)        # load\n",
    "    if GRAY_SCALE:\n",
    "        img_np = rgb2gray(img_pil)\n",
    "\n",
    "    # step 1: 先乘上H，变模糊\n",
    "    blurred = blur(img_np, blur_type)  # blur, and the line below adds noise (也就是说要先H，后再加噪声！！！！)\n",
    "    \n",
    "    # step 2: 再加上高斯噪声：\n",
    "    blurred = np.clip(blurred + np.random.normal(scale=noise_sigma/255., size=blurred.shape), 0, 1).astype(np.float32)\n",
    "    \n",
    "    #blurred = np.clip(blurred + np.random.poisson(noise_sigma/255, size=blurred.shape), 0, 1).astype(np.float32) # lam>=0 值越小，噪声频率就越少\n",
    "    \n",
    "    data_dict = { ORIGINAL: Data(img_np), \n",
    "                 CORRUPTED: Data(blurred, compare_PSNR(img_np, blurred, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)) }\n",
    "    if plot:\n",
    "        plot_dict(data_dict)\n",
    "    return data_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " DIP_TEMP 3000/3000 Loss 0.008377 psnrs: net: 20.48 x-u: 21.96 avg: 21.01\n",
      " img_number: 1.00 psnr_best_x-u: 25.100 iteration_best_temp: 427.000 \n",
      "\n",
      " img_number: 1.00 psnr_best_avg: 24.940 iteration_best_temp: 384.000 \n",
      "\n",
      " finally done! the avg_psnr of our method:  25.10031544021863\n"
     ]
    }
   ],
   "source": [
    "# ATTENTION:全部使用绝对路径！\n",
    "# SET1: 设置加入的噪声程度\n",
    "#SIGMA = 50     # small:10, medium:25, heavy:50\n",
    "\n",
    "# SET3: 选择denoiser:\n",
    "denoiser=FFDNet_color_yuan\n",
    "\n",
    "# SET4: 选择数据集：\n",
    "testset_name = 'CBSD68_mini' # set3c \n",
    "# SET5: 数据集上一层路径：\n",
    "testsets = '/home/yuanweimin/PHD_3/YUAN_LASTEST_WORK/all_data/testsets/'     # fixed\n",
    "\n",
    "# 算法名字：\n",
    "method_name='YUAN_6_18'\n",
    "# 获取数据集的全部路径：\n",
    "L_path = os.path.join(testsets, testset_name) # L_path, for Low-quality images\n",
    "\n",
    "# 保存结果的路径：\n",
    "results='/home/yuanweimin/PHD_3/YUAN_LASTEST_WORK/results/deblur/gaussion/15/'\n",
    "\n",
    "# 保存结果的路径下的保存结果文件：\n",
    "result_name=testset_name+'_'+method_name\n",
    "E_path = os.path.join(results,result_name)\n",
    "# 生成保存结果的文件夹：\n",
    "util.mkdir(E_path)\n",
    "\n",
    "# 获取数据集下所有的路径\n",
    "L_paths = util.get_image_paths(L_path)\n",
    "#print(L_paths[0])\n",
    "\n",
    "name=DIP_TEMP\n",
    "psnr_best_total=[]\n",
    "img_best_total=[]\n",
    "\n",
    "for idx,img in enumerate(L_paths): # idx:序号， img:图片\n",
    "    img_name, ext = os.path.splitext(os.path.basename(img))\n",
    "\n",
    "    # Get the LR and HR images\n",
    "    data_dict = load_imgs_deblurring(L_paths[idx], BLUR_TYPE, NOISE_SIGMA, plot=False)\n",
    "    net, net_input = get_network_and_input(img_shape=data_dict[ORIGINAL].img.shape)\n",
    "    H = get_h(data_dict[CORRUPTED].img.shape[0], BLUR_TYPE, USE_FOURIER, dtype)\n",
    "    plot_checkpoints = {1, 10, 100, 1000, 1500} # no use\n",
    "    \n",
    "    clean, psnr_hr,psnr_net_list,image_list,psnr_avg_list = train_via_admm(net, net_input, denoiser, H, data_dict[CORRUPTED].img,\n",
    "                           algorithm_name=name, plot_array=plot_checkpoints,\n",
    "                           clean_img=data_dict[ORIGINAL].img)\n",
    "    data_dict[name] = Data(clean, compare_PSNR(data_dict[ORIGINAL].img, clean, on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE))\n",
    "    # 获取psnr_avg_list列表的最大值和对应索引：\n",
    "    max1=max(psnr_hr)                       \n",
    "    index1=psnr_hr.index(max1)   \n",
    "    print('\\n img_number: %.2f psnr_best_x-u: %.3f iteration_best_temp: %.3f ' %((idx+1),max1,index1))\n",
    "    # 获取psnr_avg_list列表的最大值和对应索引：\n",
    "    max2=max(psnr_avg_list)                       \n",
    "    index2=psnr_avg_list.index(max2)   \n",
    "    print('\\n img_number: %.2f psnr_best_avg: %.3f iteration_best_temp: %.3f ' %((idx+1),max2,index2))\n",
    "\n",
    "    # # 获取最好的图像：\n",
    "    optimal_img=image_list[index1]\n",
    "    psnr_best_total.append(max1)\n",
    "    img_best_total.append(optimal_img)\n",
    "    # #save best image \n",
    "    np.clip(optimal_img, 0, 1, out=optimal_img) \n",
    "    plt.imsave(os.path.join(E_path, img_name+ext),(optimal_img).transpose(1,2,0))\n",
    "\n",
    "# 输出在相应的数据集上的avg_psnr:\n",
    "print('\\n finally done! the avg_psnr of our method: ', sum(psnr_best_total)/len(psnr_best_total))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7 (default, Sep 16 2021, 13:09:58) \n[GCC 7.5.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9befbe484a3ee3e6bab1911a773be5f36935b4d7cfa315003ecd3b6e957c7a01"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}