{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":36073,"databundleVersionId":4341859,"sourceType":"competition"}],"dockerImageVersionId":30177,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"![](https://storage.googleapis.com/kaggle-competitions/kaggle/36073/logos/header.png?t=2022-04-12-20-53-30)","metadata":{"execution":{"iopub.status.busy":"2022-04-21T18:39:54.474251Z","iopub.execute_input":"2022-04-21T18:39:54.474986Z","iopub.status.idle":"2022-04-21T18:39:55.157888Z","shell.execute_reply.started":"2022-04-21T18:39:54.474943Z","shell.execute_reply":"2022-04-21T18:39:55.157081Z"}}},{"cell_type":"markdown","source":"# The task is to denoise “shabby” pages, reversing the effects imposed on them by Augraphy","metadata":{}},{"cell_type":"markdown","source":"### Importing some basic libraries required to play with the images","metadata":{}},{"cell_type":"code","source":"from glob import glob\nimport numpy as np\nfrom matplotlib import pylab as plt\nimport cv2\nimport tensorflow as tf\nprint(tf.__version__)\nimport tensorflow.keras.layers as layers\nfrom tensorflow.keras.models import Model\nimport time\nimport os\nimport os\nimport cv2\nimport glob\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-11-18T16:15:46.211954Z","iopub.execute_input":"2024-11-18T16:15:46.212700Z","iopub.status.idle":"2024-11-18T16:15:46.219545Z","shell.execute_reply.started":"2024-11-18T16:15:46.212658Z","shell.execute_reply":"2024-11-18T16:15:46.218753Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Project Root directory:\nBASE_DIR = \"/kaggle/input/denoising-shabby-pages\"","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:46.221121Z","iopub.execute_input":"2024-11-18T16:15:46.221475Z","iopub.status.idle":"2024-11-18T16:15:46.231703Z","shell.execute_reply.started":"2024-11-18T16:15:46.221436Z","shell.execute_reply":"2024-11-18T16:15:46.230916Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"glob.glob(BASE_DIR+'/*')","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:46.233098Z","iopub.execute_input":"2024-11-18T16:15:46.233305Z","iopub.status.idle":"2024-11-18T16:15:46.246262Z","shell.execute_reply.started":"2024-11-18T16:15:46.233279Z","shell.execute_reply":"2024-11-18T16:15:46.245474Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cleaned_img_paths = sorted(glob.glob(BASE_DIR+\"/train/train/train_cleaned/*.png\"))\ntrain_shabby_paths = sorted(glob.glob(BASE_DIR+\"/train/train/train_shabby/*.png\"))","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:46.247337Z","iopub.execute_input":"2024-11-18T16:15:46.247592Z","iopub.status.idle":"2024-11-18T16:15:46.264116Z","shell.execute_reply.started":"2024-11-18T16:15:46.247537Z","shell.execute_reply":"2024-11-18T16:15:46.263579Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(cleaned_img_paths),len(train_shabby_paths)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:46.265614Z","iopub.execute_input":"2024-11-18T16:15:46.265817Z","iopub.status.idle":"2024-11-18T16:15:46.270795Z","shell.execute_reply.started":"2024-11-18T16:15:46.265792Z","shell.execute_reply":"2024-11-18T16:15:46.270104Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Display some of the shabby images given for training","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(20, 5))\nprint(\"Input samples:\")\nfor i,pth in enumerate(train_shabby_paths[0:20]):\n    im = cv2.imread(pth)\n    plt.subplot(2, 10, i+1)\n    plt.imshow(im, cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:46.271848Z","iopub.execute_input":"2024-11-18T16:15:46.272048Z","iopub.status.idle":"2024-11-18T16:15:48.089065Z","shell.execute_reply.started":"2024-11-18T16:15:46.272020Z","shell.execute_reply":"2024-11-18T16:15:48.088341Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.figure(figsize=(20, 5))\nprint(\"Target samples:\")\nfor i,pth in enumerate(cleaned_img_paths[0:20]):\n    im = cv2.imread(pth)\n    plt.subplot(2, 10, i+1)\n    plt.imshow(im, cmap='gray')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:48.090298Z","iopub.execute_input":"2024-11-18T16:15:48.090954Z","iopub.status.idle":"2024-11-18T16:15:49.853557Z","shell.execute_reply.started":"2024-11-18T16:15:48.090916Z","shell.execute_reply":"2024-11-18T16:15:49.852789Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport matplotlib.pyplot as plt\nfrom PIL import Image, ImageFilter\nimport numpy as np\nfrom PIL import ImageFont, ImageDraw, ImageOps\nfrom random import randint\nimport matplotlib.font_manager as fm\nimport pandas as pd\nfrom IPython import display\n#from utils import Logger\n\nimport torch\nfrom torch import nn, optim\nfrom torch.autograd.variable import Variable\nfrom torchvision import transforms, datasets\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nimport cv2","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:49.854749Z","iopub.execute_input":"2024-11-18T16:15:49.854988Z","iopub.status.idle":"2024-11-18T16:15:49.862856Z","shell.execute_reply.started":"2024-11-18T16:15:49.854956Z","shell.execute_reply":"2024-11-18T16:15:49.861971Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"clean_image_path = BASE_DIR+'/train/train/train_cleaned/'\nblur_image_path = BASE_DIR+'/train/train/train_shabby/'\n\nclean_image = np.asarray(cv2.imread(clean_image_path + '0001-USPS-dmm300_608.pdf-15.png',0))\nblur_image = np.asarray(cv2.imread(blur_image_path + '0001-USPS-dmm300_608.pdf-15.png',0))\n\n_, ax = plt.subplots(1,2, figsize=(20,10))\nax[0].imshow(clean_image, cmap='gray')\nax[0].title.set_text('Clean Image')\n\nax[1].imshow(blur_image, cmap='gray')\nax[1].title.set_text('Blur Image')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:49.863898Z","iopub.execute_input":"2024-11-18T16:15:49.864099Z","iopub.status.idle":"2024-11-18T16:15:50.386976Z","shell.execute_reply.started":"2024-11-18T16:15:49.864074Z","shell.execute_reply":"2024-11-18T16:15:50.386197Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(clean_image.shape, blur_image.shape)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.388014Z","iopub.execute_input":"2024-11-18T16:15:50.388239Z","iopub.status.idle":"2024-11-18T16:15:50.393245Z","shell.execute_reply.started":"2024-11-18T16:15:50.388208Z","shell.execute_reply":"2024-11-18T16:15:50.392467Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Function to get patches in each page\ndef getPatches(blur_image,clean_image,mystride):\n    blure_patches=[]\n    clean_patches=[]\n    #Defining the dimension of each patch\n    h =  ((blur_image.shape [0] // 256) +1)*256 \n    w =  ((blur_image.shape [1] // 256 ) +1)*256\n    image_padding=np.ones((h,w))\n    #Slicing method to zoom/crop in certain part of the page\n    image_padding[:blur_image.shape[0],:blur_image.shape[1]]=blur_image\n    #\n    for j in range (0,h-256,mystride):  #128 not 64\n        for k in range (0,w-256,mystride):\n            blure_patches.append(image_padding[j:j+256,k:k+256])\n    \n    h =  ((clean_image.shape [0] // 256) +1)*256 \n    w =  ((clean_image.shape [1] // 256 ) +1)*256\n\n    image_padding=np.ones((h,w))\n    image_padding[:clean_image.shape[0],:clean_image.shape[1]]=clean_image\n\n    for j in range (0,h-256,mystride):    #128 not 64\n        for k in range (0,w-256,mystride):\n            clean_patches.append(image_padding[j:j+256,k:k+256])  \n            \n    return np.array(blure_patches),np.array(clean_patches)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.395774Z","iopub.execute_input":"2024-11-18T16:15:50.395982Z","iopub.status.idle":"2024-11-18T16:15:50.407754Z","shell.execute_reply.started":"2024-11-18T16:15:50.395957Z","shell.execute_reply":"2024-11-18T16:15:50.406880Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n%%time\n_blur_patches , _cleanimage_patches = getPatches(blur_image, clean_image, mystride=128+64)\nprint(_blur_patches.shape, _cleanimage_patches.shape)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.408857Z","iopub.execute_input":"2024-11-18T16:15:50.409136Z","iopub.status.idle":"2024-11-18T16:15:50.420000Z","shell.execute_reply.started":"2024-11-18T16:15:50.409102Z","shell.execute_reply":"2024-11-18T16:15:50.419247Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"all_files = os.listdir(BASE_DIR+\"/train/train/train_shabby\")\nlen(all_files)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.420945Z","iopub.execute_input":"2024-11-18T16:15:50.421150Z","iopub.status.idle":"2024-11-18T16:15:50.429723Z","shell.execute_reply.started":"2024-11-18T16:15:50.421125Z","shell.execute_reply":"2024-11-18T16:15:50.428976Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_patches():\n  blur_patches = []\n  clean_patches = []\n\n  for file in all_files:\n    clean_image = np.asarray(cv2.imread(clean_image_path + file,0))\n    blur_image = np.asarray(cv2.imread(blur_image_path + file,0))\n\n    _blur_patches , _clean_patches = getPatches(blur_image, clean_image, mystride=128+64)\n    blur_patches.append(_blur_patches)    \n    clean_patches.append(_clean_patches)\n\n  blur_patches = np.concatenate(blur_patches, axis=0)\n  clean_patches = np.concatenate(clean_patches, axis=0)\n\n  return blur_patches, clean_patches","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.430786Z","iopub.execute_input":"2024-11-18T16:15:50.431474Z","iopub.status.idle":"2024-11-18T16:15:50.437899Z","shell.execute_reply.started":"2024-11-18T16:15:50.431435Z","shell.execute_reply":"2024-11-18T16:15:50.437216Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"blur_patches, clean_patches = create_patches()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:15:50.439131Z","iopub.execute_input":"2024-11-18T16:15:50.439662Z","iopub.status.idle":"2024-11-18T16:16:13.319185Z","shell.execute_reply.started":"2024-11-18T16:15:50.439624Z","shell.execute_reply":"2024-11-18T16:16:13.318526Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(blur_patches.shape, clean_patches.shape)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.320206Z","iopub.execute_input":"2024-11-18T16:16:13.320446Z","iopub.status.idle":"2024-11-18T16:16:13.325199Z","shell.execute_reply.started":"2024-11-18T16:16:13.320418Z","shell.execute_reply":"2024-11-18T16:16:13.324493Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class blur(Dataset):\n\n  def __init__(self, blur_patches, clean_patches, transform=None):\n    self.clean_patches = clean_patches\n    self.blur_patches = blur_patches\n    self.transform = transform\n\n  def __len__(self):\n    return len(self.clean_patches)\n\n  def __getitem__(self, idx):\n    clean_image = self.clean_patches[idx]\n    blur_image = self.blur_patches[idx]\n\n    clean_image = np.asarray(clean_image)\n    blur_image = np.asarray(blur_image)\n\n    clean_image = Image.fromarray(clean_image.astype(np.uint8))\n    blur_image = Image.fromarray(blur_image.astype(np.uint8))\n\n    if self.transform:\n       clean_image = self.transform(clean_image)\n       blur_image = self.transform(blur_image)\n\n    return blur_image, clean_image\n\ndataset = blur(blur_patches, clean_patches, transform=transforms.Compose(\n            [transforms.ToTensor(), \n             transforms.Normalize([0.5], [0.5])]))\n\nbs = 4\ndata_loader = DataLoader(dataset, batch_size=bs, shuffle=True)\nnum_batches = len(data_loader)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.326278Z","iopub.execute_input":"2024-11-18T16:16:13.326516Z","iopub.status.idle":"2024-11-18T16:16:13.592691Z","shell.execute_reply.started":"2024-11-18T16:16:13.326489Z","shell.execute_reply":"2024-11-18T16:16:13.591879Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"x,y = next(iter(data_loader))\n_, ax = plt.subplots(1,2)\nax[0].imshow(x[0].squeeze(), cmap='gray')\nax[1].imshow(y[0].squeeze(), cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.593827Z","iopub.execute_input":"2024-11-18T16:16:13.594096Z","iopub.status.idle":"2024-11-18T16:16:13.961444Z","shell.execute_reply.started":"2024-11-18T16:16:13.594059Z","shell.execute_reply":"2024-11-18T16:16:13.960795Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(x.shape, y.shape)\n","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.962405Z","iopub.execute_input":"2024-11-18T16:16:13.962593Z","iopub.status.idle":"2024-11-18T16:16:13.967000Z","shell.execute_reply.started":"2024-11-18T16:16:13.962568Z","shell.execute_reply":"2024-11-18T16:16:13.966294Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DoubleConv, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass GeneratorNet(nn.Module):\n    def __init__(\n            self, in_channels=3, out_channels=1, features=[64, 128, 256, 512],\n    ):\n        super(GeneratorNet, self).__init__()\n        self.ups = nn.ModuleList()\n        self.downs = nn.ModuleList()\n        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)\n\n        # Down part of UNET :  Refer to pdf file in architecture of generator to understand more about U-net\n        for feature in features:\n            self.downs.append(DoubleConv(in_channels, feature))\n            in_channels = feature\n\n        # Up part of UNET: Refer to pdf file in architecture of generator to understand more about U-net\n        for feature in reversed(features):\n            self.ups.append(\n                nn.ConvTranspose2d(\n                    feature*2, feature, kernel_size=2, stride=2,\n                )\n            )\n            self.ups.append(DoubleConv(feature*2, feature))\n\n        self.bottleneck = DoubleConv(features[-1], features[-1]*2)\n        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)\n\n    def forward(self, x):\n        skip_connections = []\n\n        for down in self.downs:\n            x = down(x)\n            skip_connections.append(x)\n            x = self.pool(x)\n\n        x = self.bottleneck(x)\n        skip_connections = skip_connections[::-1]\n\n        for idx in range(0, len(self.ups), 2):\n            x = self.ups[idx](x)\n            skip_connection = skip_connections[idx//2]\n\n            if x.shape != skip_connection.shape:\n                x = TF.resize(x, size=skip_connection.shape[2:])\n\n            concat_skip = torch.cat((skip_connection, x), dim=1)\n            x = self.ups[idx+1](concat_skip)\n\n        return self.final_conv(x)\n    \ndef images_to_vectors(images):\n    return images.view(images.size(0), 784)\n\ndef vectors_to_images(vectors):\n    return vectors.view(vectors.size(0), 1, 256, 256)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.968106Z","iopub.execute_input":"2024-11-18T16:16:13.968310Z","iopub.status.idle":"2024-11-18T16:16:13.986302Z","shell.execute_reply.started":"2024-11-18T16:16:13.968284Z","shell.execute_reply":"2024-11-18T16:16:13.985600Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DiscriminatorNet(nn.Module):\n    def __init__(self, in_channels):\n        super().__init__()\n        self.disc = nn.Sequential(            \n            nn.Conv2d(in_channels, 64, 3, 1, 1), # layer 1 64x256x256\n            nn.LeakyReLU(0.01, inplace=True),                     \n            nn.MaxPool2d(2),                     # layer 2 64x128x128\n            nn.Conv2d(64, 128, 3, 1, 1),         # layer 3 128x128x128\n            nn.LeakyReLU(0.01, inplace=True), \n            nn.MaxPool2d(2),                     # layer 4 128x64x64\n            nn.Conv2d(128, 256, 3, 1, 1),        # layer 5 256x64x64\n            nn.LeakyReLU(0.01, inplace=True),    # batch norm next\n            nn.BatchNorm2d(256),\n            nn.MaxPool2d(2),                     # layer 6 256x32x32\n            nn.Conv2d(256, 256, 3, 1, 1),        # layer 7 256x32x32\n            nn.LeakyReLU(0.01, inplace=True),    # batch norm next\n            nn.BatchNorm2d(256),\n            nn.MaxPool2d(2),                     # layer 8 256x16x16\n            nn.Conv2d(256, 256, 3, 1, 1),        # layer 9 256x16x16\n            nn.LeakyReLU(0.01, inplace=True),    # batch norm next\n            nn.BatchNorm2d(256),\n            nn.Conv2d(256, 1, 3, 1, 1),          # layer 10 1x16x16\n        )\n\n    def forward(self, x,y):\n        input = torch.cat([x,y], axis=1)\n        return self.disc(input)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:13.987407Z","iopub.execute_input":"2024-11-18T16:16:13.987909Z","iopub.status.idle":"2024-11-18T16:16:14.000024Z","shell.execute_reply.started":"2024-11-18T16:16:13.987870Z","shell.execute_reply":"2024-11-18T16:16:13.999365Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def noise(size):\n    n = Variable(torch.randn(size, 100))\n    if torch.cuda.is_available(): return n.cuda() \n    return n","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:14.001059Z","iopub.execute_input":"2024-11-18T16:16:14.001599Z","iopub.status.idle":"2024-11-18T16:16:14.009600Z","shell.execute_reply.started":"2024-11-18T16:16:14.001561Z","shell.execute_reply":"2024-11-18T16:16:14.008922Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndiscriminator = DiscriminatorNet(in_channels=2)\ngenerator = GeneratorNet(in_channels=1, out_channels=1)\n\ndiscriminator.cuda()\ngenerator.cuda()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:14.010536Z","iopub.execute_input":"2024-11-18T16:16:14.010738Z","iopub.status.idle":"2024-11-18T16:16:19.098511Z","shell.execute_reply.started":"2024-11-18T16:16:14.010713Z","shell.execute_reply":"2024-11-18T16:16:19.097718Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Optimizers\nd_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\ng_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n#e of image. Let us take an image of size 399×399 and imagine we only have 3×3 convolutions at our disposal. How many convolution layers will it take to reach RF of 399 ? Approxi\n# Loss function\nloss1 = nn.MSELoss()\nloss2 = nn.BCELoss()\nloss3 = nn.BCEWithLogitsLoss()\n\n# Number of steps to apply to the discriminator\nd_steps = 1  \n# Number of epochs\nnum_epochs = 100\n","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:19.099805Z","iopub.execute_input":"2024-11-18T16:16:19.100084Z","iopub.status.idle":"2024-11-18T16:16:19.106972Z","shell.execute_reply.started":"2024-11-18T16:16:19.100047Z","shell.execute_reply":"2024-11-18T16:16:19.106258Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def test_noise():\n  x,y = next(iter(data_loader))\n  return Variable(x).float().cuda()","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:19.107961Z","iopub.execute_input":"2024-11-18T16:16:19.108172Z","iopub.status.idle":"2024-11-18T16:16:19.122937Z","shell.execute_reply.started":"2024-11-18T16:16:19.108146Z","shell.execute_reply":"2024-11-18T16:16:19.122309Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#logger = Logger(model_name='DE-GAN', data_name='blur')\nnum_test_samples = bs\nfrom tqdm import tqdm\n\nfor epoch in tqdm(range(num_epochs)):\n    for n_batch, (blur_batch, clean_batch) in enumerate(data_loader):\n\n        real_data = clean_batch.float().cuda()\n        noised_data = blur_batch.float().cuda()\n        \n        # 1. Train Discriminator\n        # Generate fake data\n        fake_data = generator(noised_data)\n\n        # Reset gradients\n        d_optimizer.zero_grad()\n        \n        # 1.1 Train on Real Data\n        prediction_real = discriminator(real_data, noised_data)\n\n        # Calculate error and backpropagate\n        real_data_target = torch.ones_like(prediction_real)\n        loss_real = loss1(prediction_real, real_data_target)\n\n        # 1.2 Train on Fake Data, you would need to add one more component\n        prediction_fake = discriminator(fake_data, noised_data)\n\n        # Calculate error and backpropagate\n        fake_data_target = torch.zeros_like(prediction_real)\n        loss_fake = loss1(prediction_fake, fake_data_target)\n\n        loss_d = (loss_real + loss_fake)/2\n        loss_d.backward(retain_graph=True)\n        \n        # 1.3 Update weights with gradients\n        d_optimizer.step()\n  \n        # 2. Train Generator\n        g_optimizer.zero_grad()\n\n        # Sample noise and generate fake data\n        prediction = discriminator(fake_data, real_data)\n        \n        # Calculate error and backpropagate\n        real_data_target = torch.ones_like(prediction)\n        #import pdb; pdb.set_trace();\n\n        loss_g1 = loss1(prediction, real_data_target)\n        loss_g2 = loss1(fake_data, real_data)*500\n        loss_g = loss_g1 + loss_g2\n\n        loss_g.backward()\n\n        # Update weights with gradients\n        g_optimizer.step()\n                \n            \n    print(loss_d, loss_g, epoch, n_batch, num_batches)","metadata":{"execution":{"iopub.status.busy":"2024-11-18T16:16:19.124125Z","iopub.execute_input":"2024-11-18T16:16:19.124763Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"path ='./'\ntorch.save(generator.state_dict(), path + 'generator-deblur-final.pt')\ntorch.save(discriminator.state_dict(), path + 'discriminator-deblur-final.pt')\n\nmodel = GeneratorNet(in_channels=1, out_channels=1)\nmodel.load_state_dict(torch.load(path + 'generator-deblur-final.pt'))\nmodel.cuda()\nmodel.eval()\n\ndef calculate_psnr(img1, img2, max_value=255):\n    \"\"\"\"Calculating peak signal-to-noise ratio (PSNR) between two images.\"\"\"\n    mse = np.mean((np.array(img1, dtype=np.float32) - np.array(img2, dtype=np.float32)) ** 2)\n    if mse == 0:\n        return 100\n    return 20 * np.log10(max_value / (np.sqrt(mse)))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"blur_image, clean_image = next(iter(data_loader))\noutput = model(blur_image.cuda().float())\n\npsnr = calculate_psnr(blur_image[0].detach().cpu(), output[0].detach().cpu())\nprint()\nprint('PSNR between Clean Image and Deblur Image: ', psnr)\nprint()\n_, ax = plt.subplots(1,3, figsize=(22,10))\nax[0].imshow(blur_image[0].squeeze(), cmap='gray')\nax[0].title.set_text('Blur Image')\n\nax[1].imshow(output[0].detach().cpu().squeeze(), cmap='gray')\nax[1].title.set_text('Predicted Image')\nax[2].imshow(clean_image[0].detach().cpu().squeeze(), cmap='gray')\nax[2].title.set_text('Clean Image')\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}