Feedback

README.md

@@ -0,0 +1,21 @@
+# WCT2
+This repository contains PyTorch implementation of the paper [Photorealistic Style Transfer via Wavelet Transforms](https://arxiv.org/abs/1903.09760).
+- Avinash Prabhu - 2018102027
+- Fiza Husain - 2018101035
+- Mallika Subramanian - 2018101041
+- Tanvi Karandikar - 2018101059
+
+## How to run the code
+```python
+cd src/
+python main.py
+```
+### Arguments
+- `--content`: FOLDER-PATH-TO-CONTENT-IMAGES
+- `--content_segment`: FOLDER-PATH-TO-CONTENT-SEGMENT-LABEL-IMAGES
+- `--style`: FOLDER-PATH-TO-STYLE-IMAGES
+- `--style_segment`: FOLDER-PATH-TO-STYLE-SEGMENT-LABEL-IMAGES
+- `--output`: FOLDER-PATH-TO-OUTPUT-IMAGES
+- `--image_size`: output image size
+- `--alpha`: alpha determines the blending ratio between content and stylized features
+

images/convert.py

@@ -0,0 +1,10 @@
+import os
+from PIL import Image
+
+for subdir, dirs, files in os.walk('./'):
+    for file in files:
+        if(file.endswith('.png')):
+            img = Image.open(file)
+            li = file.split('.')[0]
+            num = li[2:]
+            img.save('tar'+num+'.png')

src/encoder_decoder.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from harr import WavePool,WaveUnpool
+
+class WaveEncoder(nn.Module):
+    def __init__(self, option_unpool):
+        super(WaveEncoder, self).__init__()
+        self.option_unpool = option_unpool
+
+        self.pad = nn.ReflectionPad2d(1)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv0 = nn.Conv2d(3, 3, 1, 1, 0)
+        self.conv1_1 = nn.Conv2d(3, 64, 3, 1, 0)
+        self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 0)
+        self.pool1 = WavePool(64)
+
+        self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 0)
+        self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 0)
+        self.pool2 = WavePool(128)
+
+        self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 0)
+        self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.conv3_4 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.pool3 = WavePool(256)
+
+        self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 0)
+
+    def forward(self, x):
+        skips = {}
+        for level in [1, 2, 3, 4]:
+            x = self.encode(x, skips, level)
+        return x
+
+    def encode(self, x, skips, level):
+        assert level in {1, 2, 3, 4}
+        if self.option_unpool == 'sum':
+            if level == 1:
+                out = self.conv0(x)
+                out = self.relu(self.conv1_1(self.pad(out)))
+                out = self.relu(self.conv1_2(self.pad(out)))
+                skips['conv1_2'] = out
+                LL, LH, HL, HH = self.pool1(out)
+                skips['pool1'] = [LH, HL, HH]
+                return LL
+            elif level == 2:
+                out = self.relu(self.conv2_1(self.pad(x)))
+                out = self.relu(self.conv2_2(self.pad(out)))
+                skips['conv2_2'] = out
+                LL, LH, HL, HH = self.pool2(out)
+                skips['pool2'] = [LH, HL, HH]
+                return LL
+            elif level == 3:
+                out = self.relu(self.conv3_1(self.pad(x)))
+                out = self.relu(self.conv3_2(self.pad(out)))
+                out = self.relu(self.conv3_3(self.pad(out)))
+                out = self.relu(self.conv3_4(self.pad(out)))
+                skips['conv3_4'] = out
+                LL, LH, HL, HH = self.pool3(out)
+                skips['pool3'] = [LH, HL, HH]
+                return LL
+            else:
+                return self.relu(self.conv4_1(self.pad(x)))
+
+        elif self.option_unpool == 'cat5':
+            if level == 1:
+                out = self.conv0(x)
+                out = self.relu(self.conv1_1(self.pad(out)))
+                return out
+
+            elif level == 2:
+                out = self.relu(self.conv1_2(self.pad(x)))
+                skips['conv1_2'] = out
+                LL, LH, HL, HH = self.pool1(out)
+                skips['pool1'] = [LH, HL, HH]
+                out = self.relu(self.conv2_1(self.pad(LL)))
+                return out
+
+            elif level == 3:
+                out = self.relu(self.conv2_2(self.pad(x)))
+                skips['conv2_2'] = out
+                LL, LH, HL, HH = self.pool2(out)
+                skips['pool2'] = [LH, HL, HH]
+                out = self.relu(self.conv3_1(self.pad(LL)))
+                return out
+
+            else:
+                out = self.relu(self.conv3_2(self.pad(x)))
+                out = self.relu(self.conv3_3(self.pad(out)))
+                out = self.relu(self.conv3_4(self.pad(out)))
+                skips['conv3_4'] = out
+                LL, LH, HL, HH = self.pool3(out)
+                skips['pool3'] = [LH, HL, HH]
+                out = self.relu(self.conv4_1(self.pad(LL)))
+                return out
+        else:
+            raise NotImplementedError
+
+
+class WaveDecoder(nn.Module):
+    def __init__(self, option_unpool):
+        super(WaveDecoder, self).__init__()
+        self.option_unpool = option_unpool
+
+        if option_unpool == 'sum':
+            multiply_in = 1
+        elif option_unpool == 'cat5':
+            multiply_in = 5
+        else:
+            raise NotImplementedError
+
+        self.pad = nn.ReflectionPad2d(1)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv4_1 = nn.Conv2d(512, 256, 3, 1, 0)
+
+        self.recon_block3 = WaveUnpool(256, option_unpool)
+        if option_unpool == 'sum':
+            self.conv3_4 = nn.Conv2d(256*multiply_in, 256, 3, 1, 0)
+        else:
+            self.conv3_4_2 = nn.Conv2d(256*multiply_in, 256, 3, 1, 0)
+        self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0)
+        self.conv3_1 = nn.Conv2d(256, 128, 3, 1, 0)
+
+        self.recon_block2 = WaveUnpool(128, option_unpool)
+        if option_unpool == 'sum':
+            self.conv2_2 = nn.Conv2d(128*multiply_in, 128, 3, 1, 0)
+        else:
+            self.conv2_2_2 = nn.Conv2d(128*multiply_in, 128, 3, 1, 0)
+        self.conv2_1 = nn.Conv2d(128, 64, 3, 1, 0)
+
+        self.recon_block1 = WaveUnpool(64, option_unpool)
+        if option_unpool == 'sum':
+            self.conv1_2 = nn.Conv2d(64*multiply_in, 64, 3, 1, 0)
+        else:
+            self.conv1_2_2 = nn.Conv2d(64*multiply_in, 64, 3, 1, 0)
+        self.conv1_1 = nn.Conv2d(64, 3, 3, 1, 0)
+
+    def forward(self, x, skips):
+        for level in [4, 3, 2, 1]:
+            x = self.decode(x, skips, level)
+        return x
+
+    def decode(self, x, skips, level):
+        assert level in {4, 3, 2, 1}
+        if level == 4:
+            out = self.relu(self.conv4_1(self.pad(x)))
+            LH, HL, HH = skips['pool3']
+            original = skips['conv3_4'] if 'conv3_4' in skips.keys() else None
+            out = self.recon_block3(out, LH, HL, HH, original)
+            _conv3_4 = self.conv3_4 if self.option_unpool == 'sum' else self.conv3_4_2
+            out = self.relu(_conv3_4(self.pad(out)))
+            out = self.relu(self.conv3_3(self.pad(out)))
+            return self.relu(self.conv3_2(self.pad(out)))
+        elif level == 3:
+            out = self.relu(self.conv3_1(self.pad(x)))
+            LH, HL, HH = skips['pool2']
+            original = skips['conv2_2'] if 'conv2_2' in skips.keys() else None
+            out = self.recon_block2(out, LH, HL, HH, original)
+            _conv2_2 = self.conv2_2 if self.option_unpool == 'sum' else self.conv2_2_2
+            return self.relu(_conv2_2(self.pad(out)))
+        elif level == 2:
+            out = self.relu(self.conv2_1(self.pad(x)))
+            LH, HL, HH = skips['pool1']
+            original = skips['conv1_2'] if 'conv1_2' in skips.keys() else None
+            out = self.recon_block1(out, LH, HL, HH, original)
+            _conv1_2 = self.conv1_2 if self.option_unpool == 'sum' else self.conv1_2_2
+            return self.relu(_conv1_2(self.pad(out)))
+        else:
+            return self.conv1_1(self.pad(x))

src/harr.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+def get_harr_wav(in_channels, pool=True):
+
+    """wavelet decomposition using conv2d
+    Defining the 4 harr wavelet kernels which when 
+    used for pooling results in 4 channels. Dim of the
+    four kerns are al 2x2 kernel matrices.
+
+    """
+    harr_wav_L = 1 / np.sqrt(2) * np.ones((1, 2))
+    harr_wav_H = 1 / np.sqrt(2) * np.ones((1, 2))
+    harr_wav_H[0][0] = - harr_wav_H[0][0]
+
+    harr_wav_LL = (harr_wav_L).T * harr_wav_L
+    harr_wav_LH = (harr_wav_L).T * harr_wav_H
+    harr_wav_HL = (harr_wav_H).T * harr_wav_L
+    harr_wav_HH = (harr_wav_H).T * harr_wav_H
+
+    #prepare the layers for the convolution
+    filter_LL = torch.from_numpy(harr_wav_LL).unsqueeze(0)
+    filter_LH = torch.from_numpy(harr_wav_LH).unsqueeze(0)
+    filter_HL = torch.from_numpy(harr_wav_HL).unsqueeze(0)
+    filter_HH = torch.from_numpy(harr_wav_HH).unsqueeze(0)
+
+    #if it is pooling using wavelet transform - use convolution
+    #if it is unpooling using wavelet transform - use component wise transposed convolution
+    if pool:
+        net = nn.Conv2d
+    else:
+        net = nn.ConvTranspose2d
+
+    #define the layers such that each channel is convolved individually for each of the 4 kernels    
+    LL = net(in_channels, in_channels,
+             kernel_size=2, stride=2, padding=0, bias=False,
+             groups=in_channels)
+    LH = net(in_channels, in_channels,
+             kernel_size=2, stride=2, padding=0, bias=False,
+             groups=in_channels)
+    HL = net(in_channels, in_channels,
+             kernel_size=2, stride=2, padding=0, bias=False,
+             groups=in_channels)
+    HH = net(in_channels, in_channels,
+             kernel_size=2, stride=2, padding=0, bias=False,
+             groups=in_channels)
+
+    LL.weight.requires_grad = False
+    LH.weight.requires_grad = False
+    HL.weight.requires_grad = False
+    HH.weight.requires_grad = False
+
+    LL.weight.data = filter_LL.float().unsqueeze(0).expand(in_channels, -1, -1, -1)
+    LH.weight.data = filter_LH.float().unsqueeze(0).expand(in_channels, -1, -1, -1)
+    HL.weight.data = filter_HL.float().unsqueeze(0).expand(in_channels, -1, -1, -1)
+    HH.weight.data = filter_HH.float().unsqueeze(0).expand(in_channels, -1, -1, -1)
+
+    return LL, LH, HL, HH
+
+class WavePool(nn.Module):
+    def __init__(self, in_channels):
+        super(WavePool, self).__init__()
+        self.LL, self.LH, self.HL, self.HH = get_harr_wav(in_channels)
+
+    def forward(self, x):
+        return self.LL(x), self.LH(x), self.HL(x), self.HH(x)
+
+
+class WaveUnpool(nn.Module):
+    def __init__(self, in_channels, option_unpool='cat5'):
+        super(WaveUnpool, self).__init__()
+        self.in_channels = in_channels
+        self.option_unpool = option_unpool
+        self.LL, self.LH, self.HL, self.HH = get_harr_wav(self.in_channels, pool=False)
+
+    def forward(self, LL, LH, HL, HH, original=None):
+        if self.option_unpool == 'sum':
+            return self.LL(LL) + self.LH(LH) + self.HL(HL) + self.HH(HH)
+        elif self.option_unpool == 'cat5' and original is not None:
+            return torch.cat([self.LL(LL), self.LH(LH), self.HL(HL), self.HH(HH), original], dim=1)
+        else:
+            raise NotImplementedError