stylegan, esrgan, srgan code

Author: aladdinpersson

ML/Pytorch/GANs/ESRGAN/ESRGAN.png

@@ -0,0 +1,48 @@
+import torch
+from PIL import Image
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+LOAD_MODEL = True
+SAVE_MODEL = True
+CHECKPOINT_GEN = "gen.pth"
+CHECKPOINT_DISC = "disc.pth"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+LEARNING_RATE = 1e-4
+NUM_EPOCHS = 10000
+BATCH_SIZE = 16
+LAMBDA_GP = 10
+NUM_WORKERS = 4
+HIGH_RES = 128
+LOW_RES = HIGH_RES // 4
+IMG_CHANNELS = 3
+
+highres_transform = A.Compose(
+    [
+        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1]),
+        ToTensorV2(),
+    ]
+)
+
+lowres_transform = A.Compose(
+    [
+        A.Resize(width=LOW_RES, height=LOW_RES, interpolation=Image.BICUBIC),
+        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1]),
+        ToTensorV2(),
+    ]
+)
+
+both_transforms = A.Compose(
+    [
+        A.RandomCrop(width=HIGH_RES, height=HIGH_RES),
+        A.HorizontalFlip(p=0.5),
+        A.RandomRotate90(p=0.5),
+    ]
+)
+
+test_transform = A.Compose(
+    [
+        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1]),
+        ToTensorV2(),
+    ]
+)

ML/Pytorch/GANs/ESRGAN/ESRGAN_generator.pth

@@ -0,0 +1,49 @@
+import torch
+from tqdm import tqdm
+import time
+import torch.nn
+import os
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import config
+from PIL import Image
+import cv2
+
+
+class MyImageFolder(Dataset):
+    def __init__(self, root_dir):
+        super(MyImageFolder, self).__init__()
+        self.data = []
+        self.root_dir = root_dir
+        self.class_names = os.listdir(root_dir)
+
+        for index, name in enumerate(self.class_names):
+            files = os.listdir(os.path.join(root_dir, name))
+            self.data += list(zip(files, [index] * len(files)))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        img_file, label = self.data[index]
+        root_and_dir = os.path.join(self.root_dir, self.class_names[label])
+
+        image = cv2.imread(os.path.join(root_and_dir, img_file))
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        both_transform = config.both_transforms(image=image)["image"]
+        low_res = config.lowres_transform(image=both_transform)["image"]
+        high_res = config.highres_transform(image=both_transform)["image"]
+        return low_res, high_res
+
+
+def test():
+    dataset = MyImageFolder(root_dir="data/")
+    loader = DataLoader(dataset, batch_size=8)
+
+    for low_res, high_res in loader:
+        print(low_res.shape)
+        print(high_res.shape)
+
+
+if __name__ == "__main__":
+    test()

ML/Pytorch/GANs/ESRGAN/config.py

@@ -0,0 +1,19 @@
+import torch.nn as nn
+from torchvision.models import vgg19
+import config
+
+
+class VGGLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.vgg = vgg19(pretrained=True).features[:35].eval().to(config.DEVICE)
+
+        for param in self.vgg.parameters():
+            param.requires_grad = False
+
+        self.loss = nn.MSELoss()
+
+    def forward(self, input, target):
+        vgg_input_features = self.vgg(input)
+        vgg_target_features = self.vgg(target)
+        return self.loss(vgg_input_features, vgg_target_features)

ML/Pytorch/GANs/ESRGAN/dataset.py

@@ -0,0 +1,154 @@
+import torch
+from torch import nn
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, use_act, **kwargs):
+        super().__init__()
+        self.cnn = nn.Conv2d(
+            in_channels,
+            out_channels,
+            **kwargs,
+            bias=True,
+        )
+        self.act = nn.LeakyReLU(0.2, inplace=True) if use_act else nn.Identity()
+
+    def forward(self, x):
+        return self.act(self.cnn(x))
+
+
+class UpsampleBlock(nn.Module):
+    def __init__(self, in_c, scale_factor=2):
+        super().__init__()
+        self.upsample = nn.Upsample(scale_factor=scale_factor, mode="nearest")
+        self.conv = nn.Conv2d(in_c, in_c, 3, 1, 1, bias=True)
+        self.act = nn.LeakyReLU(0.2, inplace=True)
+
+    def forward(self, x):
+        return self.act(self.conv(self.upsample(x)))
+
+
+class DenseResidualBlock(nn.Module):
+    def __init__(self, in_channels, channels=32, residual_beta=0.2):
+        super().__init__()
+        self.residual_beta = residual_beta
+        self.blocks = nn.ModuleList()
+
+        for i in range(5):
+            self.blocks.append(
+                ConvBlock(
+                    in_channels + channels * i,
+                    channels if i <= 3 else in_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                    use_act=True if i <= 3 else False,
+                )
+            )
+
+    def forward(self, x):
+        new_inputs = x
+        for block in self.blocks:
+            out = block(new_inputs)
+            new_inputs = torch.cat([new_inputs, out], dim=1)
+        return self.residual_beta * out + x
+
+
+class RRDB(nn.Module):
+    def __init__(self, in_channels, residual_beta=0.2):
+        super().__init__()
+        self.residual_beta = residual_beta
+        self.rrdb = nn.Sequential(*[DenseResidualBlock(in_channels) for _ in range(3)])
+
+    def forward(self, x):
+        return self.rrdb(x) * self.residual_beta + x
+
+
+class Generator(nn.Module):
+    def __init__(self, in_channels=3, num_channels=64, num_blocks=23):
+        super().__init__()
+        self.initial = nn.Conv2d(
+            in_channels,
+            num_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=True,
+        )
+        self.residuals = nn.Sequential(*[RRDB(num_channels) for _ in range(num_blocks)])
+        self.conv = nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1)
+        self.upsamples = nn.Sequential(
+            UpsampleBlock(num_channels), UpsampleBlock(num_channels),
+        )
+        self.final = nn.Sequential(
+            nn.Conv2d(num_channels, num_channels, 3, 1, 1, bias=True),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(num_channels, in_channels, 3, 1, 1, bias=True),
+        )
+
+    def forward(self, x):
+        initial = self.initial(x)
+        x = self.conv(self.residuals(initial)) + initial
+        x = self.upsamples(x)
+        return self.final(x)
+
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channels=3, features=[64, 64, 128, 128, 256, 256, 512, 512]):
+        super().__init__()
+        blocks = []
+        for idx, feature in enumerate(features):
+            blocks.append(
+                ConvBlock(
+                    in_channels,
+                    feature,
+                    kernel_size=3,
+                    stride=1 + idx % 2,
+                    padding=1,
+                    use_act=True,
+                ),
+            )
+            in_channels = feature
+
+        self.blocks = nn.Sequential(*blocks)
+        self.classifier = nn.Sequential(
+            nn.AdaptiveAvgPool2d((6, 6)),
+            nn.Flatten(),
+            nn.Linear(512 * 6 * 6, 1024),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(1024, 1),
+        )
+
+    def forward(self, x):
+        x = self.blocks(x)
+        return self.classifier(x)
+
+def initialize_weights(model, scale=0.1):
+    for m in model.modules():
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight.data)
+            m.weight.data *= scale
+
+        elif isinstance(m, nn.Linear):
+            nn.init.kaiming_normal_(m.weight.data)
+            m.weight.data *= scale
+
+
+def test():
+    gen = Generator()
+    disc = Discriminator()
+    low_res = 24
+    x = torch.randn((5, 3, low_res, low_res))
+    gen_out = gen(x)
+    disc_out = disc(gen_out)
+
+    print(gen_out.shape)
+    print(disc_out.shape)
+
+if __name__ == "__main__":
+    test()
+
+
+
+
+