Add SRGAN PyTorch Implementation (pclubiitk#35)

* Add SRGAN PyTorch Implementation

* Formatted README to reduce clutter

* Update README code block

* Fix typo

* Add new directory

* Change directory

Author: namangup

super_resolution/SRGAN_PyTorch/README.md

@@ -0,0 +1,105 @@
+# PyTorch Implementation of SRGAN
+
+## Usage
+```bash
+$ python3 main.py
+```
+following are the arguments:
+```
+usage: main.py [-h] [--root_dir ROOT_DIR] [--num_workers NUM_WORKERS] [--batch_size BATCH_SIZE] [--num_epochs NUM_EPOCHS] [--lr LR]
+               [--pre_num_epochs PRE_NUM_EPOCHS] [--outdir OUTDIR] [--load_checkpoint LOAD_CHECKPOINT] [--b B]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --root_dir ROOT_DIR   path to dataset
+  --num_workers NUM_WORKERS
+                        number of data loading workers
+  --batch_size BATCH_SIZE
+                        input batch size
+  --num_epochs NUM_EPOCHS
+                        number of epochs to train for
+  --lr LR               learning rate
+  --pre_num_epochs PRE_NUM_EPOCHS
+                        number of pre-training epochs
+  --outdir OUTDIR       directory to output model checkpoints
+  --load_checkpoint LOAD_CHECKPOINT
+                        Pass 1 to load checkpoint
+  --b B                 number of residual blocks in generator
+```
+## Contributed by:
+[Naman Gupta](https://github.com/namangup)
+
+## References
+* **Title**: Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network
+* **Authors**: Christian Ledig et. al.
+* **Link**: https://arxiv.org/abs/1609.04802
+* **Tags**: Super Resolution, Generative Adversarial Networks
+* **Year**: 2017
+
+## What's New?
+Super Resolution received substantial attention from within the computer
+vision research community and has a wide range of
+applications. The optimization target of supervised
+SR algorithms is commonly the minimization of the mean
+squared error (MSE) between the recovered HR image
+and the ground truth. This is convenient as minimizing
+MSE also maximizes the peak signal-to-noise ratio (PSNR),
+which is a common measure used to evaluate and compare
+SR algorithms. However, the ability of MSE (and
+PSNR) to capture perceptually relevant differences, such
+as high texture detail, is very limited as they are defined
+based on pixel-wise image difference. Hence, to capure those details, SRGAN define a novel perceptual loss using high-level feature maps of the VGG network
+combined with a discriminator that encourages solutions
+perceptually hard to distinguish from the HR reference
+images.
+## Architecture
+![model](assets/model.jpg)
+## Loss Functions
+we have the following loss functions:\
+**Perceptual Loss**\
+![perceptual loss](assets/perceptual_loss.png)\
+**PixelWise MSE Loss**\
+![MSE loss](assets/MSE_loss.png)\
+**Content Loss**\
+![content loss](assets/content_loss.png)\
+**Adversarial Loss**\
+![adversarial loss](assets/adversarial_loss.png)\
+
+## Implementation
+Following the paper the SRResNet(Generator) is pre-trained first on MSE Loss, followed by adversarial training of both the Generator and Discriminator alternately (k=1).
+For the perceptual loss, VGG(5,4) is used by default.
+
+The dataset consists of ~40k images randomly sampled from the Imagenet Dataset, and 96\*96 patches are cropped randomly from each image. 
+These 96\*96 images are downsampled and fed into the generator as images of size 24\*24, which in turn generates back Super Resolution images of size 96\*96.
+
+## Results
+
+On pre-training the Generator for 100 epochs followed by adversarial training for 200 epochs, the following results are obtained.
+
+>***NOTE*** : Go to the assets folder to view full sized images. They have been resized for better readability.
+
+>*x4 refers to image upscaled four times*.
+
+**Low Resolution (Original)** 
+<p float="left">
+<img src="assets/lr.png" width="280" style="margin:10px">
+</p>
+
+**x4 Bicubic Interpolation, High Resolution (Original)**
+<p float="left">
+<img src="assets/lr_bicubic.png" width="400" style="margin:10px">
+<img src="assets/hr.png" width="400" style="margin:10px">
+</p>
+
+**x4 *Online Image Enhancer*, x4 SRGAN**
+<p>
+<img src="assets/lr_letsenhance.png" width="400" style="margin:10px">
+<img src="assets/sr.png" width="400" style="margin:10px">
+</p>
+
+>I used [letsenhance.io](https://letsenhance.io/) which claims to use a "*Powerful AI to increase image resolution without quality loss*".
+
+The SRGAN generated image clearly retains more features, and produces better images.
+
+The current model can definitely achieve far better results given more data and training epochs . 
+I used 40k images whereas, the authors used 350k images, trained for 10^5 steps.

super_resolution/SRGAN_PyTorch/assets/MSE_loss.png

@@ -0,0 +1,36 @@
+from torch.utils.data import Dataset
+from PIL import Image
+import os
+from torchvision import transforms
+
+def clean_dataset(dir):
+    """ Remove images which are not in RGB colour space"""
+    for img in os.listdir(dir):
+        path = os.path.join(dir, img)
+        im = Image.open(path)
+        if(im.mode != 'RGB'):
+            os.remove(path)
+
+class TrainDataset(Dataset):
+
+    def __init__(self, dir):
+        super().__init__()
+        clean_dataset(dir)
+        self.img = [os.path.join(dir, x) for x in os.listdir(dir)]
+        self.hr = transforms.Compose([
+                                    transforms.RandomCrop(96, pad_if_needed=True),
+                                    transforms.ToTensor(),
+        ])
+        self.lr = transforms.Compose([
+                                    transforms.ToPILImage(),
+                                    transforms.Resize(24, interpolation=Image.BICUBIC),
+                                    transforms.ToTensor()
+        ])
+    
+    def __getitem__(self, index):
+        hr_image = self.hr(Image.open(self.img[index]))
+        lr_image = self.lr(hr_image)
+        return lr_image, hr_image
+
+    def __len__(self):
+        return len(self.img)

super_resolution/SRGAN_PyTorch/assets/adversarial_loss.png

@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import models, utils
+from torch.utils.data import DataLoader
+import time
+from dataloader import TrainDataset
+from models import FeatureExtractor, Generator, Discriminator
+from torchsummary import summary
+import argparse
+import os
+import matplotlib.pyplot as plt
+from torch.utils.tensorboard import SummaryWriter
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--root_dir', default='./', help='path to dataset')
+parser.add_argument('--num_workers', type=int, default=2, help='number of data loading workers')
+parser.add_argument('--batch_size', type=int, default=128, help='input batch size')
+parser.add_argument('--num_epochs', type=int, default=200, help='number of epochs to train for')
+parser.add_argument('--lr', type=float, default=0.0001, help='learning rate')
+parser.add_argument('--pre_num_epochs', type=int, default=100, help='number of pre-training epochs')
+parser.add_argument('--outdir', default='./', help='directory to output model checkpoints')
+parser.add_argument('--load_checkpoint', default=0, type=int, help='Pass 1 to load checkpoint')
+parser.add_argument('--b', default=16, type=int, help='number of residual blocks in generator')
+args = parser.parse_args()
+
+# Load data
+dataset = TrainDataset(args.root_dir)
+dataloader = DataLoader(dataset, args.batch_size, True, num_workers=args.num_workers)
+# Initialize models
+vgg = models.vgg19(pretrained=True)
+feature_extractor = FeatureExtractor(vgg, 5, 4)
+if torch.cuda.device_count() > 1:
+    feature_extractor = nn.DataParallel(feature_extractor)
+feature_extractor = feature_extractor.to(device)
+
+disc = Discriminator()
+if torch.cuda.device_count() > 1:
+    disc = nn.DataParallel(disc)
+disc = disc.to(device)
+if args.load_checkpoint == 1 and os.path.exists('disc.pt'):
+    disc.load_state_dict(torch.load('disc.pt'))
+print(disc)
+
+gen = Generator(args.b)
+if torch.cuda.device_count() > 1:
+    gen = nn.DataParallel(gen)
+gen = gen.to(device)
+if args.load_checkpoint == 1 and os.path.exists('gen.pt'):
+    gen.load_state_dict(torch.load('gen.pt'))
+print(gen)
+
+content_criterion = nn.MSELoss()
+adversarial_criterion = nn.BCELoss()
+optimG = optim.Adam(gen.parameters(), args.lr)
+schedulerG1 = optim.lr_scheduler.MultiStepLR(optimG, [100], 0.1)
+schedulerG2 = optim.lr_scheduler.MultiStepLR(optimG, [100], 0.1)
+optimD = optim.Adam(disc.parameters(), args.lr)
+schedulerD = optim.lr_scheduler.MultiStepLR(optimD, [100], 0.1)
+writer = SummaryWriter()
+# Generator pre-training
+start_time = time.time()
+iters = 0
+for epoch in range(args.pre_num_epochs):
+    
+    for i, data in enumerate(dataloader, 0):
+
+        lr, hr_real = data
+        hr_real = hr_real.to(device)
+        lr = lr.to(device)
+
+        batch_size = hr_real.size()[0]
+        hr_fake = gen(lr)
+
+        gen.zero_grad()
+        gen_content_loss = content_criterion(hr_fake, hr_real)
+        gen_content_loss.backward()
+        optimG.step()
+
+        if i == 0:
+            print(f'[{epoch}/{args.pre_num_epochs}][{i}/{len(dataloader)}] Gen_MSE: {gen_content_loss.item()}')
+        iters += 1 
+
+    torch.save(gen.state_dict(), f'{args.outdir}gen.pt')
+    schedulerG1.step()
+    print(f'Time Elapsed: {(time.time()-start_time): .2f}')
+
+# Adversarial Training
+G_losses = []
+D_losses = []
+iters = 0
+optimG = optim.Adam(gen.parameters(), args.lr)
+for epoch in range(args.num_epochs):
+    
+    for i, data in enumerate(dataloader):
+        iters += 1
+        lr, hr_real = data
+        batch_size = hr_real.size()[0]
+        hr_real = hr_real.to(device)
+        lr = lr.to(device)
+        hr_fake = gen(lr)
+
+        # Label Smoothing (Salimans et. al. 2016)
+        target_real = torch.rand(batch_size, 1, device=device)*0.85+0.3
+        target_fake = torch.rand(batch_size, 1, device=device)*0.15
+
+        # Discriminator
+        disc.zero_grad()
+        D_x = disc(hr_real)
+        D_G_z1 = disc(hr_fake.detach())
+        errD_real = adversarial_criterion(D_x, target_real)
+        errD_fake = adversarial_criterion(D_G_z1, target_fake)
+        errD = errD_real + errD_fake
+        D_x = D_x.view(-1).mean().item()
+        D_G_z1 = D_G_z1.view(-1).mean().item()
+        errD.backward()
+        optimD.step()
+
+        # Generator
+        gen.zero_grad()
+        real_features = feature_extractor(hr_real)
+        fake_features = feature_extractor(hr_fake)
+        ones = torch.ones(batch_size, 1, device=device)
+
+        errG_mse = content_criterion(hr_fake, hr_real)
+        errG_vgg = content_criterion(fake_features, real_features)
+        D_G_z2 = disc(hr_fake)
+        errG_adv = adversarial_criterion(D_G_z2, ones)
+        errG = errG_mse + 0.006*errG_vgg + 0.001*errG_adv
+        D_G_z2 = D_G_z2.view(-1).mean().item()
+        errG.backward()
+        optimG.step()
+        if i == 0:
+            print(f'[{epoch}/{args.num_epochs}][{i}/{len(dataloader)}] errD: {errD.item():.4f}'
+                    f' errG: {errG.item():.4f} ({errG_mse.item():.4f}/{0.006*errG_vgg.item():.4f}/{0.001*errG_adv.item():.4f})'
+            f' D(HR): {D_x :.4f} D(G(LR1)): {D_G_z1:.4f} D(G(LR2)): {D_G_z2:.4f}')
+        
+        G_losses.append(errG.item())
+        D_losses.append(errD.item())
+
+    torch.save(gen.state_dict(), f'{args.outdir}gen.pt')
+    torch.save(disc.state_dict(), f'{args.outdir}disc.pt')
+    print(f'Time Elapsed: {(time.time()-start_time): .2f}')
+    schedulerD.step()
+    schedulerG2.step()
+
+print(f'Finished Training {args.num_epochs} epochs')
+
+plt.figure(figsize=(10,5))
+plt.title("Generator and Discriminator Loss During Training")
+plt.plot(G_losses,label="G")
+plt.plot(D_losses,label="D")
+plt.xlabel("Iterations")
+plt.ylabel("Loss")
+plt.legend()
+plt.show()