ADMM_DIP_WTV_denoising.ipynb

{"cells":[{"cell_type":"markdown","metadata":{"id":"hkhmUkEoK9sN"},"source":["This is the code of our **ADMM DIP-WTV** described in the paper:\n","\n","https://arxiv.org/abs/2009.11380\n","\n","- Select `name` below to switch between the images.\n","- Change the `weight`update to exploit the UPEN principle.\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"ztgYHO8HLHtT"},"outputs":[],"source":["# Mounting my Google Drive and set the cwd on my imagefolder\n","from google.colab import drive\n","drive.mount('/content/drive/', force_remount=True)\n","%cd /content/drive/My Drive/ADMM-DIPTV"]},{"cell_type":"markdown","metadata":{"id":"MNjQBbjxK9sQ"},"source":["# Import libs"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"XlkW1Jp3K9sR"},"outputs":[],"source":["from __future__ import print_function\n","import matplotlib.pyplot as plt\n","%matplotlib inline\n","\n","import os\n","\n","import numpy as np\n","from models import *\n","\n","import torch\n","import torch.optim\n","\n","from skimage.metrics import peak_signal_noise_ratio as compare_psnr\n","from utils.denoising_utils import *\n","from utils.utils import *\n","import random\n","\n","torch.backends.cudnn.enabled = True\n","torch.backends.cudnn.benchmark =True\n","dtype = torch.cuda.FloatTensor\n","\n","imsize = -1\n","PLOT = True\n","sigma = 20 #70\n","sigma_ = sigma/255.\n","\n","torch.manual_seed(42)\n","np.random.seed(42)\n","random.seed(42)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"7-oWpWwtK9sb"},"outputs":[],"source":["name = 'butterfly.png'\n","fname = f\"data2/denoising/%s\"%(name)"]},{"cell_type":"markdown","metadata":{"id":"Q9jGUSHpK9si"},"source":["# Load image"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Xp9n4_P2K9sk"},"outputs":[],"source":["# Add synthetic noise\n","img_pil = crop_image(get_image(fname, imsize)[0], d=32)\n","img_np = pil_to_np(img_pil)  \n","    \n","# pay attention!\n","if img_np.shape[0]==2:\n","  img_np_temp = np.zeros((3,img_np.shape[1], img_np.shape[2]))\n","  img_np_temp[0,:,:] = img_np[0,:,:]  \n","  img_np_temp[1,:,:] = img_np[0,:,:]\n","  img_np_temp[2,:,:] = img_np[0,:,:]\n","  img_np = img_np_temp\n","if img_np.shape[0]==4:\n","  img_np = img_np[:-1,:,:]\n","    \n","img_noisy_pil, img_noisy_np = get_noisy_image(img_np, sigma_)\n","\n","if PLOT:\n","    plot_image_grid([img_np, img_noisy_np], 4, 5);\n","\n","print(img_np.shape)"]},{"cell_type":"markdown","metadata":{"id":"ImXzJDYlK9sr"},"source":["# Setup"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"fM3raZF0K9st"},"outputs":[],"source":["INPUT = 'noise' \n","pad = 'reflection'\n","OPT_OVER = 'net' \n","\n","reg_noise_std = 1./20 \n","\n","show_every = 5\n","\n","num_iter = 100 \n","input_depth = 3 \n","figsize = 1 \n","    \n","    \n","net = get_net(input_depth, 'skip', pad,\n","              skip_n33d=128, \n","              skip_n33u=128, \n","              skip_n11=4, \n","              num_scales=5,\n","              upsample_mode='bilinear').type(dtype)\n","    \n","net_input = get_noise(input_depth, INPUT, (img_pil.size[1], img_pil.size[0])).type(dtype).detach()\n","\n","# Compute number of parameters\n","s  = sum([np.prod(list(p.size())) for p in net.parameters()]); \n","print ('Number of params: %d' % s)\n","\n","# Loss\n","mse = torch.nn.MSELoss().type(dtype)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"YffHWcVes8Nz"},"outputs":[],"source":["size = img_np.shape\n","h = size[-2]\n","w = size[-1]\n","Dh_psf = np.array([ [0, 0, 0], [1, -1, 0], [0, 0, 0]])\n","Dv_psf = np.array([ [0, 1, 0], [0, -1, 0], [0, 0, 0]])\n","Id_psf = np.array([[1]])\n","\n","Id_DFT = torch.from_numpy(psf2otf(Id_psf, [h,w])).cuda()\n","Dh_DFT = torch.from_numpy(psf2otf(Dh_psf, [h,w])).cuda()\n","Dv_DFT = torch.from_numpy(psf2otf(Dv_psf, [h,w])).cuda()\n","\n","DhT_DFT = torch.conj(Dh_DFT)\n","DvT_DFT = torch.conj(Dv_DFT)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"MTeTGsNmhwVb"},"outputs":[],"source":["LR = 1e-3\n","optimizer = torch.optim.Adam(net.parameters(), lr=LR)"]},{"cell_type":"markdown","metadata":{"id":"tsYDjymLK9s2"},"source":["# Optimize"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"36E_XzpiK9s3","scrolled":true},"outputs":[],"source":["net_input_saved = net_input.detach().clone()\n","noise = net_input.detach().clone()\n","out_avg = None\n","last_net = None\n","psrn_noisy_last = 0\n","\n","img_noisy_torch = np_to_torch(img_noisy_np).type(dtype) \n","u = 0*img_noisy_torch.detach().clone()\n","t_h = 0*img_noisy_torch.detach().clone()\n","t_v = 0*img_noisy_torch.detach().clone()\n","\n","mu_f = torch.zeros(net_input.shape, device=0)\n","mu_t_h = torch.zeros(net_input.shape, device=0)\n","mu_t_v = torch.zeros(net_input.shape, device=0)\n","\n","beta_t = 10\n","inner_iterations = 50\n","\n","loss_values = []\n","fun_values = []\n","psnr_values = []\n","running_loss=0\n","\n","for i in range(num_iter):\n","    \n","    if inner_iterations>1:\n","      optimizer = torch.optim.Adam(net.parameters(), lr=LR)\n","\n","    for j in range(inner_iterations):\n","      optimizer.zero_grad()\n","    \n","      #First problem\n","      out = net(net_input)\n","      [Dh_out, Dv_out] = D(out, Dh_DFT, Dv_DFT)\n","\n","      total_loss = norm2_loss(out-img_noisy_torch)\n","      total_loss += (beta_t/2)*norm2_loss(Dh_out-(t_h-mu_t_h).detach()) + (beta_t/2)*norm2_loss(Dv_out-(t_v-mu_t_v).detach())\n","      \n","      total_loss.backward()\n","      optimizer.step()\n","      \n","    running_loss = total_loss.item()\n","    loss_values.append(running_loss)\n","\n","    out = net(net_input)\n","  \n","    [Dh_out, Dv_out] = D(out, Dh_DFT, Dv_DFT)\n","\n","    #TV problem: second problem \n","    q_h                 = Dh_out + mu_t_h\n","    q_v                 = Dv_out + mu_t_v\n","    q_norm              = torch.sqrt(torch.pow(q_h,2) + torch.pow(q_v,2))\n","    weight              = torch.div(torch.pow(torch.norm(out-img_noisy_torch),2)/(6*h*w),q_norm)\n","    weight              = weight.detach().clone()\n","    q_norm[q_norm == 0] = weight[q_norm == 0]/beta_t\n","    q_norm              = torch.clamp( q_norm - weight/beta_t , min=0 )/q_norm\n","    t_h                 = (q_norm * q_h).detach().clone()\n","    t_v                 = (q_norm * q_v).detach().clone()\n","\n","    #Ascent step: updating lagrangian parameter\n","    mu_t_h = (mu_t_h + (Dh_out- t_h)).detach().clone()\n","    mu_t_v = (mu_t_v + (Dv_out- t_v)).detach().clone()\n","\n","    psrn_noisy = compare_psnr(img_noisy_np, out.detach().cpu().numpy()[0]) \n","    psrn_gt    = compare_psnr(img_np, out.detach().cpu().numpy()[0]) \n","    psnr_values.append(psrn_gt)\n","    \n","    if  PLOT and i % show_every == 0:\n","        out_np = torch_to_np(out)\n","        plot_image_grid([np.clip(out_np, 0, 1), img_np], factor=5)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"SKXelO3iK9tI"},"outputs":[],"source":["out_np = torch_to_np(net(net_input))\n","plot_image_grid([np.clip(out_np, 0, 1), img_np,img_noisy_np], factor=10);"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Dpky5iA92T8N"},"outputs":[],"source":["plt.plot(psnr_values)\n","plt.title(f\"Best PSNR:%.2f - Last PSNR:%.2f\"%(np.max(psnr_values), psnr_values[-1]))\n","plt.show()\n","\n","plt.plot(loss_values)\n","plt.title(\"Running loss\")\n","plt.show()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"FGN-zYzHiVRM"},"outputs":[],"source":["from PIL import Image\n","\n","#Saving output image\n","\n","folder = \"ADMM_DIPWTV\"\n","path = f\"results/%s/%s_%s_%d\"%(folder, folder, name[:-4], sigma)\n","\n","out_np=255*out_np\n","print(np.amax(out_np))\n","u=out_np.transpose(1, 2, 0)\n","np.savetxt(f'%s.txt'%path,u[:,:,0])\n","\n","u=u.astype(np.uint8)\n","print(u.shape)\n","\n","im = Image.fromarray(u)\n","im.save(f'%s.png'%path)\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"7asLfXz1K-bh"},"outputs":[],"source":["#Saving output PSNR\n","folder = \"ADMM_DIPWTV\"\n","file_name = \"ADMM_DIPWTV_psnr\"\n","path = f\"results/%s/%s_%s_%d\"%(folder, file_name, name[:-4], sigma)\n","np.savetxt(f'%s.txt'%path,psnr_values)\n","\n","#Saving noisy image\n","folder = \"ADMM_DIPWTV\"\n","file_name = \"ADMM_DIPWTV_noisy\"\n","path = f\"results/%s/%s_%s_%d\"%(folder, file_name, name[:-4], sigma)\n","\n","img_noisy_np=255*img_noisy_np\n","u=img_noisy_np.transpose(1, 2, 0)\n","\n","u=u.astype(np.uint8)\n","print(u.shape)\n","\n","im = Image.fromarray(u)\n","im.save(f'%s.png'%path)"]}],"metadata":{"accelerator":"GPU","colab":{"collapsed_sections":[],"name":"ADMM_DIP_WTV_denoising.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}