denoising_block_discrepancyDRED.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6rQlzjabmuZI"
   },
   "source": [
    "# **Denoising constrained DRED PGDA**\n",
    "\n",
    "---\n",
    "\n",
    "This code is mainly based on DeepRED code available at https://github.com/GaryMataev/DeepRED\n",
    "\n",
    "This notebook is the implementation of the PGDA version of the following paper: \n",
    "\n",
    "**Constrained and unconstrained deep image prior optimization models with automatic regularization** by *Pasquale Cascarano, Giorgia Franchini, Erich Kobler, Federica Porta and Andrea Sebastiani*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "whIFzTBGmuZX"
   },
   "source": [
    "# Import libs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 1022,
     "status": "ok",
     "timestamp": 1643021209006,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "zfs4eRwEmuZZ"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from threading import Thread  # for running the denoiser in parallel\n",
    "import queue\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim\n",
    "from models.skip import skip  # our network\n",
    "\n",
    "from utils.utils import *  # auxiliary functions\n",
    "from utils.data import Data  # class that holds img, psnr, time\n",
    "\n",
    "from skimage.restoration import denoise_nl_means\n",
    "\n",
    "from scipy.signal import convolve2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 26,
     "status": "ok",
     "timestamp": 1643021209010,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "ngzxpNsXmuZb"
   },
   "outputs": [],
   "source": [
    "# got GPU? - if you are not getting the exact article results set CUDNN to False\n",
    "CUDA_FLAG = True\n",
    "CUDNN = True \n",
    "if CUDA_FLAG:\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "    # GPU accelerated functionality for common operations in deep neural nets\n",
    "    torch.backends.cudnn.enabled = CUDNN\n",
    "    # benchmark mode is good whenever your input sizes for your network do not vary.\n",
    "    # This way, cudnn will look for the optimal set of algorithms for that particular \n",
    "    # configuration (which takes some time). This usually leads to faster runtime.\n",
    "    # But if your input sizes changes at each iteration, then cudnn will benchmark every\n",
    "    # time a new size appears, possibly leading to worse runtime performances.\n",
    "    torch.backends.cudnn.benchmark = CUDNN\n",
    "    # torch.backends.cudnn.deterministic = True\n",
    "    dtype = torch.cuda.FloatTensor\n",
    "else:\n",
    "    dtype = torch.FloatTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qn-zwXGPmuZd"
   },
   "source": [
    "# CONSTANCTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 24,
     "status": "ok",
     "timestamp": 1643021209012,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "lMmxGWHomuZe"
   },
   "outputs": [],
   "source": [
    "SIGMA = 35\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",
    "# graphs labels:\n",
    "X_LABELS = ['Iterations']*3\n",
    "Y_LABELS = ['PSNR between x and net (db)', 'PSNR with original image (db)', 'loss']\n",
    "\n",
    "# Algorithm NAMES (to get the relevant image: use data_dict[alg_name].img)\n",
    "# for example use data_dict['Clean'].img to get the clean image\n",
    "ORIGINAL = 'Clean'\n",
    "CORRUPTED = 'Noisy'\n",
    "NLM = 'NLM'\n",
    "DIP_NLM = 'CDRED (NLM)'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TrKukM43muZf"
   },
   "source": [
    "# Load image for Denoising"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 23,
     "status": "ok",
     "timestamp": 1643021209014,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "ouHYgV4CmuZg"
   },
   "outputs": [],
   "source": [
    "def load_image(fclean, fnoisy=None, sigma=25, plot=False):\n",
    "    \"\"\" \n",
    "        fname - input file name\n",
    "        d - Make dimensions divisible by `d`\n",
    "        sigma - the amount of noise you want to add noise to the image\n",
    "        Return a numpy image, and a noisy numpy image with sigma selected\n",
    "    \"\"\"\n",
    "    _, img_np = load_and_crop_image(fclean)\n",
    "    if fnoisy is None:\n",
    "        img_noisy_np = np.clip(img_np + np.random.normal(scale=sigma / 255., size=img_np.shape), 0, 1).astype(\n",
    "            np.float32)\n",
    "        # img_noisy_np = pil_to_np(np_to_pil(img_noisy_np)) # making it an image then loading it back to numpy\n",
    "    else:\n",
    "        _, img_noisy_np = load_and_crop_image(fnoisy)\n",
    "    data_dict = {ORIGINAL: Data(img_np), CORRUPTED: Data(img_noisy_np, compare_PSNR(img_np, img_noisy_np,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": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 976
    },
    "executionInfo": {
     "elapsed": 3676,
     "status": "ok",
     "timestamp": 1643021212669,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "hnDJbDN3muZj",
    "outputId": "4f60d3a4-7239-4998-c0d6-6ef06f1032e0"
   },
   "outputs": [],
   "source": [
    "# load the image and add noise - for real use send same image file to fclean and fnoisy and ignore psnrs\n",
    "data_dict = load_image('datasets/Set5/woman_GT.bmp', sigma=SIGMA, plot=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fLtpUtx-7TBr"
   },
   "source": [
    "#  ESTIMATING THE NOISE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 64,
     "status": "ok",
     "timestamp": 1643021212673,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "dFVMghZq7WUb",
    "outputId": "71c3c310-d986-4418-8a98-f71b955c09a2"
   },
   "outputs": [],
   "source": [
    "lap_kernel = np.array([[1,-2,1], [-2, 4, -2], [1,-2,1]])\n",
    "h=data_dict[CORRUPTED].img[:,:,:].shape[2]\n",
    "w=data_dict[CORRUPTED].img[:,:,:].shape[1]\n",
    "\n",
    "def estimate_variance(img):\n",
    "  out = convolve2d(img, lap_kernel, mode='valid')\n",
    "  out = np.sum(np.abs(out))\n",
    "  out = (out*np.sqrt(0.5*np.pi)/(6*(h-2)*(w-2)))\n",
    "  return out\n",
    "\n",
    "print(data_dict[CORRUPTED].img[:,:,:].shape)\n",
    "NOISE_SIGMA = estimate_variance(data_dict[CORRUPTED].img[0,:,:])*255\n",
    "print(NOISE_SIGMA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5teVFdjUmuZk"
   },
   "source": [
    "# THE NETWORK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 41,
     "status": "ok",
     "timestamp": 1643021212674,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "qaBXQfjTmuZl"
   },
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jlEN_0DFmuZm"
   },
   "source": [
    "## The Non Local Means denoiser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 42,
     "status": "ok",
     "timestamp": 1643021212676,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "WuTyJ_ZZmuZn"
   },
   "outputs": [],
   "source": [
    "def non_local_means(noisy_np_img, sigma, fast_mode=True):\n",
    "    \"\"\" get a numpy noisy image\n",
    "        returns a denoised numpy image using Non-Local-Means\n",
    "    \"\"\" \n",
    "    sigma = sigma / 255.\n",
    "    h = 0.6 * sigma if fast_mode else 0.8 * sigma\n",
    "    patch_kw = dict(h=h,                   # Cut-off distance, a higher h results in a smoother image\n",
    "                    sigma=sigma,           # sigma provided\n",
    "                    fast_mode=fast_mode,   # If True, a fast version is used. If False, the original version is used.\n",
    "                    patch_size=5,          # 5x5 patches (Size of patches used for denoising.)\n",
    "                    patch_distance=6,      # 13x13 search area\n",
    "                    multichannel=False)\n",
    "    denoised_img = []\n",
    "    n_channels = noisy_np_img.shape[0]\n",
    "    for c in range(n_channels):\n",
    "        denoise_fast = denoise_nl_means(noisy_np_img[c, :, :], **patch_kw)\n",
    "        denoised_img += [denoise_fast]\n",
    "    return np.array(denoised_img, dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 709,
     "output_embedded_package_id": "1bP6AsuCy9hTCgzwcldvhWDa9PuVwERU6"
    },
    "executionInfo": {
     "elapsed": 3071,
     "status": "ok",
     "timestamp": 1643021215706,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "C9sOEq-PmuZp",
    "outputId": "f4238687-456e-409e-81b9-d97fcffce854"
   },
   "outputs": [],
   "source": [
    "# Run Non-Local-Means\n",
    "denoised_img = non_local_means(data_dict[CORRUPTED].img, sigma=NOISE_SIGMA)\n",
    "data_dict[NLM] = Data(denoised_img, compare_PSNR(data_dict[ORIGINAL].img, denoised_img,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE))\n",
    "plot_dict(data_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H4lrQq72muZq"
   },
   "source": [
    "# Constrained Deep Image prior via PGDA with RED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 20,
     "status": "ok",
     "timestamp": 1643021215711,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "sxvDmRxkmuZr"
   },
   "outputs": [],
   "source": [
    "def train_via_admm(net, net_input, denoiser_function, y, noise_lev,tau, org_img=None,                      # y is the noisy image\n",
    "                   plot_array={}, algorithm_name=\"\", admm_iter=5000, save_path=\"\",           # path to save params\n",
    "                   LR=0.001,                                                                      # learning rate\n",
    "                   sigma_f=3, update_iter=10, method='fixed_point',   # method: 'fixed_point' or 'grad' or 'mixed'\n",
    "                   beta=1, mu=0.5,mu_r=1, LR_x=None, noise_factor=0.033,    #0.033 LR_x needed only if method!=fixed_point\n",
    "                   threshold=40, threshold_step=0.01, increase_reg=0.03):                # increase regularization \n",
    "    \"\"\" training the network using\n",
    "        ## Must Params ##\n",
    "        net                 - the network to be trained\n",
    "        net_input           - the network input\n",
    "        denoiser_function   - an external denoiser function, used as black box, this function\n",
    "                              must get numpy noisy image, and return numpy denoised image\n",
    "        y                   - the noisy image\n",
    "        sigma               - the noise level (int 0-255)\n",
    "        \n",
    "        # optional params #\n",
    "        org_img             - the original image if exist for psnr compare only, or None (default)\n",
    "        plot_array          - prints params at the begging of the training and plot images at the required indices\n",
    "        admm_iter           - total number of admm epoch\n",
    "        LR                  - the lr of the network in admm (step 2)\n",
    "        sigma_f             - the sigma to send the denoiser function\n",
    "        update_iter         - denoised image updated every 'update_iter' iteration\n",
    "        method              - 'fixed_point' or 'grad' or 'mixed' \n",
    "        algorithm_name      - the name that would show up while running, just to know what we are running ;)\n",
    "                \n",
    "        # equation params #  \n",
    "        beta                - regularization parameter (lambda in the article)\n",
    "        mu                  - ADMM parameter\n",
    "        LR_x                - learning rate of the parameter x, needed only if method!=fixed point\n",
    "        # more\n",
    "        noise_factor       - the amount of noise added to the input of the network\n",
    "        threshold          - when the image become close to the noisy image at this psnr\n",
    "        increase_reg       - we going to increase regularization by this amount\n",
    "        threshold_step     - and keep increasing it every step\n",
    "    \"\"\"\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",
    "    if org_img is not None:\n",
    "        psnr_y = compare_PSNR(org_img, y,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)  # get the noisy image psnr\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",
    "    # optimizer and scheduler\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=LR)  # using ADAM opt\n",
    "    \n",
    "    y_torch = np_to_torch(y).type(dtype)\n",
    "\n",
    "    x, u,u_r, r, out_np, res,out_np_old = y.copy(), np.zeros_like(y),np.zeros_like(y),y.copy(), np.zeros_like(y),y.copy(),np.zeros_like(y)\n",
    "    f_x, avg, avg2, avg3 = x.copy(), np.rint(y), np.rint(y), np.rint(y)\n",
    "    list_psnr=[]\n",
    "    list_stopping=[]\n",
    "    img_queue = queue.Queue()\n",
    "    denoiser_thread = Thread(target=lambda q, f, f_args: q.put(f(*f_args)),\n",
    "                             args=(img_queue, denoiser_function, [x.copy(), sigma_f]))\n",
    "    denoiser_thread.start()\n",
    "\n",
    "    for i in range(1, 1 + admm_iter):\n",
    "\n",
    "        rho = tau*noise_lev*np.sqrt(y.shape[0]*y.shape[1]*y.shape[2] - 1)  \n",
    "\n",
    "        # step 1, update x using a denoiser and result from step 1\n",
    "        if i % update_iter == 0:  # the denoiser work in parallel\n",
    "            denoiser_thread.join()\n",
    "            f_x = img_queue.get()\n",
    "            denoiser_thread = Thread(target=lambda q, f, f_args: q.put(f(*f_args)),\n",
    "                                     args=(img_queue, denoiser_function, [x.copy(), sigma_f]))\n",
    "            denoiser_thread.start()\n",
    "\n",
    "        x_old = 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",
    "\n",
    "        # step 2 projection\n",
    "        r_old = r\n",
    "        r = res + u_r  \n",
    "        r_norm = np.sqrt(np.sum(np.square(r)))\n",
    "        if r_norm > rho:\n",
    "          r = r*(rho/r_norm)\n",
    "          r_norm_2 = np.sqrt(np.sum(np.square(r)))\n",
    "\n",
    "        # step 3, 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",
    "\n",
    "        # loss:\n",
    "        loss_y = mse(out-y_torch, np_to_torch(r_old-u_r).type(dtype))\n",
    "        loss_x = mse(out, np_to_torch(x_old - u).type(dtype))\n",
    "        total_loss = mu_r*loss_y + mu * loss_x\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        out = net(net_input)\n",
    "        out_np = torch_to_np(out)\n",
    "\n",
    "        res = out_np-y\n",
    "        \n",
    "        # step 4, update u and u_r\n",
    "        u = u +  0.001*(out_np - x)\n",
    "        u_r = u_r + 0.001*(res - r)\n",
    "\n",
    "        # Averaging:\n",
    "        avg = avg * .99 + out_np * .01\n",
    "\n",
    "        stopping = np.sqrt(np.sum(np.square(out_np-y)))/ rho \n",
    "        #stopping = 255*(np.sqrt(np.sum(np.square(org_img-y)))/ np.sqrt(y.shape[0]*y.shape[1]*y.shape[2] - 1))\n",
    "        list_stopping.append(stopping)\n",
    "\n",
    "        #if stopping < 1:\n",
    "        #  break\n",
    "        \n",
    "        out_np_old = out_np\n",
    "        \n",
    "\n",
    "        # show psnrs:\n",
    "        psnr_noisy = compare_PSNR(out_np, y)\n",
    "        if psnr_noisy > threshold:\n",
    "            mu = mu + increase_reg\n",
    "            beta = beta + increase_reg\n",
    "            threshold += threshold_step\n",
    "        if org_img is not None:\n",
    "            psnr_net, psnr_avg = compare_PSNR(org_img, out_np,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE), compare_PSNR(org_img, avg,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "            psnr_x, psnr_x_u = compare_PSNR(org_img, x,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE), compare_PSNR(org_img, x - u,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE)\n",
    "            list_psnr.append(psnr_avg)\n",
    "            print('\\r', algorithm_name, '%04d/%04d Loss %f' % (i, admm_iter, total_loss.item()),\n",
    "                  'psnrs: y: %.2f psnr_noisy: %.2f net: %.2f' % (psnr_y, psnr_noisy, psnr_net),\n",
    "                  'x: %.2f x-u: %.2f avg: %.2f' % (psnr_x, psnr_x_u, psnr_avg), 'params: rho: %.2f r_norm: %.2f  r_norm2: %.2f stopping: %.2f mu: %.2f' % (rho,r_norm,r_norm_2,stopping,mu), end='')\n",
    "            if i in plot_array:  # plot images\n",
    "                tmp_dict = {'Clean': Data(org_img),\n",
    "                            #'Noisy': Data(y, psnr_y),\n",
    "                            'Net': Data(out_np, psnr_net),\n",
    "                            #'x-u': Data(x - u, psnr_x_u),\n",
    "                            'avg': Data(avg, psnr_avg),\n",
    "                            'x': Data(x, psnr_x),\n",
    "                            'u': Data((u - np.min(u)) / (np.max(u) - np.min(u))),\n",
    "                            'u_r': Data((u_r - np.min(u_r)) / (np.max(u_r) - np.min(u_r))),\n",
    "                            'r': Data(r)\n",
    "                            }\n",
    "                plot_dict(tmp_dict)\n",
    "        else:\n",
    "            print('\\r', algorithm_name, 'iteration %04d/%04d Loss %f' % (i, admm_iter, total_loss.item()), end='')\n",
    "    \n",
    "    if denoiser_thread.is_alive():\n",
    "        denoiser_thread.join()  # joining the thread\n",
    "    return avg,list_psnr,list_stopping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xCGuOYzfmuZ2"
   },
   "source": [
    "## Let's Go:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1123915,
     "status": "ok",
     "timestamp": 1643022339607,
     "user": {
      "displayName": "Andrea Sebastiani",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjzunHGnSGYxhcCchFdbTKbFFzug5WVMQk3hsnks38=s64",
      "userId": "08381018070201927882"
     },
     "user_tz": -60
    },
    "id": "JYLWve6jmuaB",
    "outputId": "8bca98d4-1dde-40ca-8299-bc01a6f68eee"
   },
   "outputs": [],
   "source": [
    "def run_and_plot(denoiser, name, plot_checkpoints={}):\n",
    "    global data_dict\n",
    "    noise_lev = NOISE_SIGMA/255\n",
    "    tau=1\n",
    "    net, net_input = get_network_and_input(img_shape=data_dict[CORRUPTED].img.shape)\n",
    "    denoised_img,list_psnr,list_stopping = train_via_admm(net, net_input, denoiser, data_dict[CORRUPTED].img,noise_lev,tau,\n",
    "                                  plot_array=plot_checkpoints, algorithm_name=name,\n",
    "                                  org_img=data_dict[ORIGINAL].img)\n",
    "    data_dict[name] = Data(denoised_img, compare_PSNR(data_dict[ORIGINAL].img, denoised_img,on_y=(not GRAY_SCALE), gray_scale=GRAY_SCALE))\n",
    "    plot_dict(data_dict)\n",
    "\n",
    "    return denoised_img,list_psnr,list_stopping\n",
    "\n",
    "\n",
    "plot_checkpoints = {1, 10, 50, 100, 250, 500, 2000, 3500, 4000, 4500, 5000,5500,6000,6500,7000,7500,8000} \n",
    "denoised_img,list_psnr,list_stopping=run_and_plot(non_local_means, DIP_NLM, plot_checkpoints)  # you may try it with different denoisers"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "denoising_block_discrepancyDRED.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.4 ('PsiDONet22')",
   "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.10.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "2e600517162b2d52f2381e9f8081dec209e7747d3e7da5b473d1484e44801cd5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}