Add application file

data/palette.txt

@@ -0,0 +1,256 @@
+0 0 0
+128 0 0
+0 128 0
+128 128 0
+0 0 128
+128 0 128
+0 128 128
+128 128 128
+64 0 0
+191 0 0
+64 128 0
+191 128 0
+64 0 128
+191 0 128
+64 128 128
+191 128 128
+0 64 0
+128 64 0
+0 191 0
+128 191 0
+0 64 128
+128 64 128
+22 22 22
+23 23 23
+24 24 24
+25 25 25
+26 26 26
+27 27 27
+28 28 28
+29 29 29
+30 30 30
+31 31 31
+32 32 32
+33 33 33
+34 34 34
+35 35 35
+36 36 36
+37 37 37
+38 38 38
+39 39 39
+40 40 40
+41 41 41
+42 42 42
+43 43 43
+44 44 44
+45 45 45
+46 46 46
+47 47 47
+48 48 48
+49 49 49
+50 50 50
+51 51 51
+52 52 52
+53 53 53
+54 54 54
+55 55 55
+56 56 56
+57 57 57
+58 58 58
+59 59 59
+60 60 60
+61 61 61
+62 62 62
+63 63 63
+64 64 64
+65 65 65
+66 66 66
+67 67 67
+68 68 68
+69 69 69
+70 70 70
+71 71 71
+72 72 72
+73 73 73
+74 74 74
+75 75 75
+76 76 76
+77 77 77
+78 78 78
+79 79 79
+80 80 80
+81 81 81
+82 82 82
+83 83 83
+84 84 84
+85 85 85
+86 86 86
+87 87 87
+88 88 88
+89 89 89
+90 90 90
+91 91 91
+92 92 92
+93 93 93
+94 94 94
+95 95 95
+96 96 96
+97 97 97
+98 98 98
+99 99 99
+100 100 100
+101 101 101
+102 102 102
+103 103 103
+104 104 104
+105 105 105
+106 106 106
+107 107 107
+108 108 108
+109 109 109
+110 110 110
+111 111 111
+112 112 112
+113 113 113
+114 114 114
+115 115 115
+116 116 116
+117 117 117
+118 118 118
+119 119 119
+120 120 120
+121 121 121
+122 122 122
+123 123 123
+124 124 124
+125 125 125
+126 126 126
+127 127 127
+128 128 128
+129 129 129
+130 130 130
+131 131 131
+132 132 132
+133 133 133
+134 134 134
+135 135 135
+136 136 136
+137 137 137
+138 138 138
+139 139 139
+140 140 140
+141 141 141
+142 142 142
+143 143 143
+144 144 144
+145 145 145
+146 146 146
+147 147 147
+148 148 148
+149 149 149
+150 150 150
+151 151 151
+152 152 152
+153 153 153
+154 154 154
+155 155 155
+156 156 156
+157 157 157
+158 158 158
+159 159 159
+160 160 160
+161 161 161
+162 162 162
+163 163 163
+164 164 164
+165 165 165
+166 166 166
+167 167 167
+168 168 168
+169 169 169
+170 170 170
+171 171 171
+172 172 172
+173 173 173
+174 174 174
+175 175 175
+176 176 176
+177 177 177
+178 178 178
+179 179 179
+180 180 180
+181 181 181
+182 182 182
+183 183 183
+184 184 184
+185 185 185
+186 186 186
+187 187 187
+188 188 188
+189 189 189
+190 190 190
+191 191 191
+192 192 192
+193 193 193
+194 194 194
+195 195 195
+196 196 196
+197 197 197
+198 198 198
+199 199 199
+200 200 200
+201 201 201
+202 202 202
+203 203 203
+204 204 204
+205 205 205
+206 206 206
+207 207 207
+208 208 208
+209 209 209
+210 210 210
+211 211 211
+212 212 212
+213 213 213
+214 214 214
+215 215 215
+216 216 216
+217 217 217
+218 218 218
+219 219 219
+220 220 220
+221 221 221
+222 222 222
+223 223 223
+224 224 224
+225 225 225
+226 226 226
+227 227 227
+228 228 228
+229 229 229
+230 230 230
+231 231 231
+232 232 232
+233 233 233
+234 234 234
+235 235 235
+236 236 236
+237 237 237
+238 238 238
+239 239 239
+240 240 240
+241 241 241
+242 242 242
+243 243 243
+244 244 244
+245 245 245
+246 246 246
+247 247 247
+248 248 248
+249 249 249
+250 250 250
+251 251 251
+252 252 252
+253 253 253
+254 254 254
+255 255 255

libs/blocks.py

@@ -0,0 +1,739 @@
+"""Network Modules
+ - encoder3: vgg encoder up to relu31
+ - decoder3: mirror decoder to encoder3
+ - encoder4: vgg encoder up to relu41
+ - decoder4: mirror decoder to encoder4
+ - encoder5: vgg encoder up to relu51
+ - styleLoss: gram matrix loss for all style layers
+ - styleLossMask: gram matrix loss for all style layers, compare between each part defined by a mask
+ - GramMatrix: compute gram matrix for one layer
+ - LossCriterion: style transfer loss that include both content & style losses
+ - LossCriterionMask: style transfer loss that include both content & style losses, use the styleLossMask
+ - VQEmbedding: codebook class for VQVAE
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .vq_functions import vq, vq_st
+from collections import OrderedDict
+
+class MetaModule(nn.Module):
+    """
+    Base class for PyTorch meta-learning modules. These modules accept an
+    additional argument `params` in their `forward` method.
+
+    Notes
+    -----
+    Objects inherited from `MetaModule` are fully compatible with PyTorch
+    modules from `torch.nn.Module`. The argument `params` is a dictionary of
+    tensors, with full support of the computation graph (for differentiation).
+    """
+    def meta_named_parameters(self, prefix='', recurse=True):
+        gen = self._named_members(
+            lambda module: module._parameters.items()
+            if isinstance(module, MetaModule) else [],
+            prefix=prefix, recurse=recurse)
+        for elem in gen:
+            yield elem
+
+    def meta_parameters(self, recurse=True):
+        for name, param in self.meta_named_parameters(recurse=recurse):
+            yield param
+
+class BatchLinear(nn.Linear, MetaModule):
+    '''A linear meta-layer that can deal with batched weight matrices and biases, as for instance output by a
+    hypernetwork.'''
+    __doc__ = nn.Linear.__doc__
+
+    def forward(self, input, params=None):
+        if params is None:
+            params = OrderedDict(self.named_parameters())
+
+        bias = params.get('bias', None)
+        weight = params['weight']
+
+        output = input.matmul(weight.permute(*[i for i in range(len(weight.shape) - 2)], -1, -2))
+        output += bias.unsqueeze(-2)
+        return output
+
+class decoder1(nn.Module):
+    def __init__(self):
+        super(decoder1,self).__init__()
+        self.reflecPad2 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+        self.conv3 = nn.Conv2d(64,3,3,1,0)
+        # 224 x 224
+
+    def forward(self,x):
+        out = self.reflecPad2(x)
+        out = self.conv3(out)
+        return out
+
+
+class decoder2(nn.Module):
+    def __init__(self):
+        super(decoder2,self).__init__()
+        # decoder
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,64,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(64,64,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        out = self.reflecPad5(x)
+        out = self.conv5(out)
+        out = self.relu5(out)
+        out = self.unpool(out)
+        out = self.reflecPad6(out)
+        out = self.conv6(out)
+        out = self.relu6(out)
+        out = self.reflecPad7(out)
+        out = self.conv7(out)
+        return out
+
+class encoder3(nn.Module):
+    def __init__(self):
+        super(encoder3,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2,return_indices = True)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2,return_indices = True)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+    def forward(self,x):
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        out = self.relu2(out)
+        out = self.reflecPad3(out)
+        out = self.conv3(out)
+        pool1 = self.relu3(out)
+        out,pool_idx = self.maxPool(pool1)
+        out = self.reflecPad4(out)
+        out = self.conv4(out)
+        out = self.relu4(out)
+        out = self.reflecPad5(out)
+        out = self.conv5(out)
+        pool2 = self.relu5(out)
+        out,pool_idx2 = self.maxPool2(pool2)
+        out = self.reflecPad6(out)
+        out = self.conv6(out)
+        out = self.relu6(out)
+        return out
+
+class decoder3(nn.Module):
+    def __init__(self):
+        super(decoder3,self).__init__()
+        # decoder
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,128,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 x 112
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(128,128,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(128,64,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(64,64,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        output = {}
+        out = self.reflecPad7(x)
+        out = self.conv7(out)
+        out = self.relu7(out)
+        out = self.unpool(out)
+        out = self.reflecPad8(out)
+        out = self.conv8(out)
+        out = self.relu8(out)
+        out = self.reflecPad9(out)
+        out = self.conv9(out)
+        out_relu9 = self.relu9(out)
+        out = self.unpool2(out_relu9)
+        out = self.reflecPad10(out)
+        out = self.conv10(out)
+        out = self.relu10(out)
+        out = self.reflecPad11(out)
+        out = self.conv11(out)
+        return out
+
+class encoder4(nn.Module):
+    def __init__(self):
+        super(encoder4,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,256,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(256,256,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(256,256,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 28 x 28
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(256,512,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+    def forward(self,x,sF=None,matrix11=None,matrix21=None,matrix31=None):
+        output = {}
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        output['r11'] = self.relu2(out)
+        out = self.reflecPad7(output['r11'])
+
+        out = self.conv3(out)
+        output['r12'] = self.relu3(out)
+
+        output['p1'] = self.maxPool(output['r12'])
+        out = self.reflecPad4(output['p1'])
+        out = self.conv4(out)
+        output['r21'] = self.relu4(out)
+        out = self.reflecPad7(output['r21'])
+
+        out = self.conv5(out)
+        output['r22'] = self.relu5(out)
+
+        output['p2'] = self.maxPool2(output['r22'])
+        out = self.reflecPad6(output['p2'])
+        out = self.conv6(out)
+        output['r31'] = self.relu6(out)
+        if(matrix31 is not None):
+            feature3,transmatrix3 = matrix31(output['r31'],sF['r31'])
+            out = self.reflecPad7(feature3)
+        else:
+            out = self.reflecPad7(output['r31'])
+        out = self.conv7(out)
+        output['r32'] = self.relu7(out)
+
+        out = self.reflecPad8(output['r32'])
+        out = self.conv8(out)
+        output['r33'] = self.relu8(out)
+
+        out = self.reflecPad9(output['r33'])
+        out = self.conv9(out)
+        output['r34'] = self.relu9(out)
+
+        output['p3'] = self.maxPool3(output['r34'])
+        out = self.reflecPad10(output['p3'])
+        out = self.conv10(out)
+        output['r41'] = self.relu10(out)
+
+        return output
+
+class decoder4(nn.Module):
+    def __init__(self):
+        super(decoder4,self).__init__()
+        # decoder
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,256,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+        # 28 x 28
+
+        self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
+        # 56 x 56
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(256,256,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(256,256,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(256,256,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad15 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv15 = nn.Conv2d(256,128,3,1,0)
+        self.relu15 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 112 x 112
+
+        self.reflecPad16 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv16 = nn.Conv2d(128,128,3,1,0)
+        self.relu16 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad17 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv17 = nn.Conv2d(128,64,3,1,0)
+        self.relu17 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
+        # 224 x 224
+
+        self.reflecPad18 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv18 = nn.Conv2d(64,64,3,1,0)
+        self.relu18 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad19 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv19 = nn.Conv2d(64,3,3,1,0)
+
+    def forward(self,x):
+        # decoder
+        out = self.reflecPad11(x)
+        out = self.conv11(out)
+        out = self.relu11(out)
+        out = self.unpool(out)
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+        out = self.reflecPad15(out)
+        out = self.conv15(out)
+        out = self.relu15(out)
+        out = self.unpool2(out)
+        out = self.reflecPad16(out)
+        out = self.conv16(out)
+        out = self.relu16(out)
+        out = self.reflecPad17(out)
+        out = self.conv17(out)
+        out = self.relu17(out)
+        out = self.unpool3(out)
+        out = self.reflecPad18(out)
+        out = self.conv18(out)
+        out = self.relu18(out)
+        out = self.reflecPad19(out)
+        out = self.conv19(out)
+        return out
+
+class encoder5(nn.Module):
+    def __init__(self):
+        super(encoder5,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,256,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(256,256,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(256,256,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 28 x 28
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(256,512,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,512,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(512,512,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(512,512,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+
+        self.maxPool4 = nn.MaxPool2d(kernel_size=2,stride=2)
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(512,512,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+
+    def forward(self,x,sF=None,contentV256=None,styleV256=None,matrix11=None,matrix21=None,matrix31=None):
+        output = {}
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        output['r11'] = self.relu2(out)
+        out = self.reflecPad7(output['r11'])
+
+        #out = self.reflecPad3(output['r11'])
+        out = self.conv3(out)
+        output['r12'] = self.relu3(out)
+
+        output['p1'] = self.maxPool(output['r12'])
+        out = self.reflecPad4(output['p1'])
+        out = self.conv4(out)
+        output['r21'] = self.relu4(out)
+        out = self.reflecPad7(output['r21'])
+
+        #out = self.reflecPad5(output['r21'])
+        out = self.conv5(out)
+        output['r22'] = self.relu5(out)
+
+        output['p2'] = self.maxPool2(output['r22'])
+        out = self.reflecPad6(output['p2'])
+        out = self.conv6(out)
+        output['r31'] = self.relu6(out)
+        if(styleV256 is not None):
+            feature = matrix31(output['r31'],sF['r31'],contentV256,styleV256)
+            out = self.reflecPad7(feature)
+        else:
+            out = self.reflecPad7(output['r31'])
+        out = self.conv7(out)
+        output['r32'] = self.relu7(out)
+
+        out = self.reflecPad8(output['r32'])
+        out = self.conv8(out)
+        output['r33'] = self.relu8(out)
+
+        out = self.reflecPad9(output['r33'])
+        out = self.conv9(out)
+        output['r34'] = self.relu9(out)
+
+        output['p3'] = self.maxPool3(output['r34'])
+        out = self.reflecPad10(output['p3'])
+        out = self.conv10(out)
+        output['r41'] = self.relu10(out)
+
+        out = self.reflecPad11(out)
+        out = self.conv11(out)
+        out = self.relu11(out)
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.maxPool4(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+        output['r51'] = out
+        return output
+
+class styleLoss(nn.Module):
+    def forward(self, input, target):
+        ib,ic,ih,iw = input.size()
+        iF = input.view(ib,ic,-1)
+        iMean = torch.mean(iF,dim=2)
+        iCov = GramMatrix()(input)
+
+        tb,tc,th,tw = target.size()
+        tF = target.view(tb,tc,-1)
+        tMean = torch.mean(tF,dim=2)
+        tCov = GramMatrix()(target)
+
+        loss = nn.MSELoss(size_average=False)(iMean,tMean) + nn.MSELoss(size_average=False)(iCov,tCov)
+        return loss/tb
+
+class GramMatrix(nn.Module):
+    def forward(self, input):
+        b, c, h, w = input.size()
+        f = input.view(b,c,h*w) # bxcx(hxw)
+        # torch.bmm(batch1, batch2, out=None)   #
+        # batch1: bxmxp, batch2: bxpxn -> bxmxn #
+        G = torch.bmm(f,f.transpose(1,2)) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc
+        return G.div_(c*h*w)
+
+class LossCriterion(nn.Module):
+    def __init__(self, style_layers, content_layers, style_weight, content_weight, 
+                 model_path = '/home/xtli/Documents/GITHUB/LinearStyleTransfer/models/'):
+        super(LossCriterion,self).__init__()
+
+        self.style_layers = style_layers
+        self.content_layers = content_layers
+        self.style_weight = style_weight
+        self.content_weight = content_weight
+
+        self.styleLosses = [styleLoss()] * len(style_layers)
+        self.contentLosses = [nn.MSELoss()] * len(content_layers)
+
+        self.vgg5 = encoder5()
+        self.vgg5.load_state_dict(torch.load(os.path.join(model_path, 'vgg_r51.pth')))
+
+        for param in self.vgg5.parameters():
+            param.requires_grad = True
+
+    def forward(self, transfer, image, content=True, style=True):
+        cF = self.vgg5(image)
+        sF = self.vgg5(image)
+        tF = self.vgg5(transfer)
+
+        losses = {}
+
+        # content loss
+        if content:
+            totalContentLoss = 0
+            for i,layer in enumerate(self.content_layers):
+                cf_i = cF[layer]
+                cf_i = cf_i.detach()
+                tf_i = tF[layer]
+                loss_i = self.contentLosses[i]
+                totalContentLoss += loss_i(tf_i,cf_i)
+            totalContentLoss = totalContentLoss * self.content_weight
+            losses['content'] = totalContentLoss
+
+        # style loss
+        if style:
+            totalStyleLoss = 0
+            for i,layer in enumerate(self.style_layers):
+                sf_i = sF[layer]
+                sf_i = sf_i.detach()
+                tf_i = tF[layer]
+                loss_i = self.styleLosses[i]
+                totalStyleLoss += loss_i(tf_i,sf_i)
+            totalStyleLoss = totalStyleLoss * self.style_weight
+            losses['style'] = totalStyleLoss
+        
+        return losses
+
+class styleLossMask(nn.Module):
+    def forward(self, input, target, mask):
+        ib,ic,ih,iw = input.size()
+        iF = input.view(ib,ic,-1)
+        tb,tc,th,tw = target.size()
+        tF = target.view(tb,tc,-1)
+
+        loss = 0
+        mb, mc, mh, mw = mask.shape
+        for i in range(mb):
+            # resize mask to have the same size of the feature
+            maski = F.interpolate(mask[i:i+1], size = (ih, iw), mode = 'nearest')
+            mask_flat = maski.view(mc, -1)
+            for j in range(mc):
+                # get features for each part
+                idx = torch.nonzero(mask_flat[j]).squeeze()
+                if len(idx.shape) == 0 or idx.shape[0] == 0:
+                    continue
+                ipart = torch.index_select(iF, 2, idx)
+                tpart = torch.index_select(tF, 2, idx)
+
+                iMean = torch.mean(ipart,dim=2)
+                iGram = torch.bmm(ipart, ipart.transpose(1,2)).div_(ic*ih*iw) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc
+
+                tMean = torch.mean(tpart,dim=2)
+                tGram = torch.bmm(tpart, tpart.transpose(1,2)).div_(tc*th*tw) # f: bxcx(hxw), f.transpose: bx(hxw)xc -> bxcxc
+
+                loss += nn.MSELoss()(iMean,tMean) + nn.MSELoss()(iGram,tGram)
+        return loss/tb
+
+class LossCriterionMask(nn.Module):
+    def __init__(self, style_layers, content_layers, style_weight, content_weight, 
+                 model_path = '/home/xtli/Documents/GITHUB/LinearStyleTransfer/models/'):
+        super(LossCriterionMask,self).__init__()
+
+        self.style_layers = style_layers
+        self.content_layers = content_layers
+        self.style_weight = style_weight
+        self.content_weight = content_weight
+
+        self.styleLosses = [styleLossMask()] * len(style_layers)
+        self.contentLosses = [nn.MSELoss()] * len(content_layers)
+
+        self.vgg5 = encoder5()
+        self.vgg5.load_state_dict(torch.load(os.path.join(model_path, 'vgg_r51.pth')))
+
+        for param in self.vgg5.parameters():
+            param.requires_grad = True
+
+    def forward(self, transfer, image, mask, content=True, style=True):
+        # mask: B, N, H, W
+        cF = self.vgg5(image)
+        sF = self.vgg5(image)
+        tF = self.vgg5(transfer)
+
+        losses = {}
+
+        # content loss
+        if content:
+            totalContentLoss = 0
+            for i,layer in enumerate(self.content_layers):
+                cf_i = cF[layer]
+                cf_i = cf_i.detach()
+                tf_i = tF[layer]
+                loss_i = self.contentLosses[i]
+                totalContentLoss += loss_i(tf_i,cf_i)
+            totalContentLoss = totalContentLoss * self.content_weight
+            losses['content'] = totalContentLoss
+
+        # style loss
+        if style:
+            totalStyleLoss = 0
+            for i,layer in enumerate(self.style_layers):
+                sf_i = sF[layer]
+                sf_i = sf_i.detach()
+                tf_i = tF[layer]
+                loss_i = self.styleLosses[i]
+                totalStyleLoss += loss_i(tf_i,sf_i, mask)
+            totalStyleLoss = totalStyleLoss * self.style_weight
+            losses['style'] = totalStyleLoss
+        
+        return losses
+
+class VQEmbedding(nn.Module):
+    def __init__(self, K, D):
+        super().__init__()
+        self.embedding = nn.Embedding(K, D)
+        self.embedding.weight.data.uniform_(-1./K, 1./K)
+
+    def forward(self, z_e_x):
+        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
+        latents = vq(z_e_x_, self.embedding.weight)
+        return latents
+
+    def straight_through(self, z_e_x, return_index=False):
+        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
+        z_q_x_, indices = vq_st(z_e_x_, self.embedding.weight.detach())
+        z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()
+
+        z_q_x_bar_flatten = torch.index_select(self.embedding.weight,
+            dim=0, index=indices)
+        z_q_x_bar_ = z_q_x_bar_flatten.view_as(z_e_x_)
+        z_q_x_bar = z_q_x_bar_.permute(0, 3, 1, 2).contiguous()
+
+        if return_index:
+            return z_q_x, z_q_x_bar, indices
+        else:
+            return z_q_x, z_q_x_bar

libs/custom_transform.py

@@ -0,0 +1,249 @@
+import torch 
+import torchvision
+from torchvision import transforms
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+import numpy as np 
+from PIL import Image, ImageFilter
+import random 
+
+class BaseTransform(object):
+    """
+    Resize and center crop. 
+    """
+    def __init__(self, res):
+        self.res = res 
+        
+    def __call__(self, index, image):
+        image = TF.resize(image, self.res, Image.BILINEAR)
+        w, h  = image.size
+        left  = int(round((w - self.res) / 2.))
+        top   = int(round((h - self.res) / 2.))
+
+        return TF.crop(image, top, left, self.res, self.res)
+
+
+class ComposeTransform(object):
+    def __init__(self, tlist):
+        self.tlist = tlist 
+    
+    def __call__(self, index, image):
+        for trans in self.tlist:
+            image = trans(index, image)
+        
+        return image 
+
+class RandomResize(object):
+    def __init__(self, rmin, rmax, N):
+        self.reslist = [random.randint(rmin, rmax) for _ in range(N)]
+
+    def __call__(self, index, image):
+        return TF.resize(image, self.reslist[index], Image.BILINEAR)
+
+class RandomCrop(object):
+    def __init__(self, res, N):
+        self.res  = res 
+        self.cons = [(np.random.uniform(0, 1), np.random.uniform(0, 1)) for _ in range(N)] 
+    
+    def __call__(self, index, image):
+        ws, hs = self.cons[index]
+        w, h   = image.size
+        left = int(round((w-self.res)*ws))
+        top  = int(round((h-self.res)*hs))
+
+        return TF.crop(image, top, left, self.res, self.res)
+
+class RandomHorizontalFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, index, image):
+        if self.plist[index.cpu()] < self.p_ref:
+            return TF.hflip(image)
+        else:
+            return image
+
+
+class TensorTransform(object):
+    def __init__(self):
+        self.to_tensor = transforms.ToTensor()
+        #self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    
+    def __call__(self, image):
+        image = self.to_tensor(image)
+        #image = self.normalize(image)
+
+        return image
+        
+        
+class RandomGaussianBlur(object):
+    def __init__(self, sigma, p, N):
+        self.min_x = sigma[0]
+        self.max_x = sigma[1]
+        self.del_p = 1 - p
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            x = self.plist[index] - self.p_ref 
+            m = (self.max_x - self.min_x) / self.del_p
+            b = self.min_x 
+            s = m * x + b 
+
+            return image.filter(ImageFilter.GaussianBlur(radius=s))
+        else:
+            return image
+
+
+class RandomGrayScale(object):
+    def __init__(self, p, N):
+        self.grayscale = transforms.RandomGrayscale(p=1.) # Deterministic (We still want flexible out_dim).
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            return self.grayscale(image)
+        else:
+            return image
+
+
+class RandomColorBrightness(object):
+    def __init__(self, x, p, N):
+        self.min_x = max(0, 1 - x)
+        self.max_x = 1 + x
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        self.rlist = [random.uniform(self.min_x, self.max_x) for _ in range(N)]
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            return TF.adjust_brightness(image, self.rlist[index])
+        else:
+            return image
+
+
+class RandomColorContrast(object):
+    def __init__(self, x, p, N):
+        self.min_x = max(0, 1 - x)
+        self.max_x = 1 + x
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        self.rlist = [random.uniform(self.min_x, self.max_x) for _ in range(N)]
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            return TF.adjust_contrast(image, self.rlist[index])
+        else:
+            return image
+
+
+class RandomColorSaturation(object):
+    def __init__(self, x, p, N):
+        self.min_x = max(0, 1 - x)
+        self.max_x = 1 + x
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        self.rlist = [random.uniform(self.min_x, self.max_x) for _ in range(N)]
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            return TF.adjust_saturation(image, self.rlist[index])
+        else:
+            return image
+
+
+class RandomColorHue(object):
+    def __init__(self, x, p, N):
+        self.min_x = -x
+        self.max_x = x
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        self.rlist = [random.uniform(self.min_x, self.max_x) for _ in range(N)]
+
+    def __call__(self, index, image):
+        if self.plist[index] < self.p_ref:
+            return TF.adjust_hue(image, self.rlist[index])
+        else:
+            return image
+
+
+class RandomVerticalFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        
+    def __call__(self, indice, image):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+
+        if len(image.size()) == 3:
+            image_t = image[I].flip([1]) 
+        else:
+            image_t = image[I].flip([2])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+
+
+class RandomHorizontalTensorFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, indice, image, is_label=False):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+        
+        if len(image.size()) == 3:
+            image_t = image[I].flip([2])
+        else:
+            image_t = image[I].flip([3])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+
+
+class RandomResizedCrop(object):
+    def __init__(self, N, res, scale=(0.5, 1.0)):
+        self.res    = res
+        self.scale  = scale 
+        self.rscale = [np.random.uniform(*scale) for _ in range(N)]
+        self.rcrop  = [(np.random.uniform(0, 1), np.random.uniform(0, 1)) for _ in range(N)]
+
+    def random_crop(self, idx, img):
+        ws, hs = self.rcrop[idx]
+        res1 = int(img.size(-1))
+        res2 = int(self.rscale[idx]*res1)
+        i1 = int(round((res1-res2)*ws))
+        j1 = int(round((res1-res2)*hs))
+
+        return img[:, :, i1:i1+res2, j1:j1+res2]
+
+
+    def __call__(self, indice, image):
+        new_image = []
+        res_tar   = self.res // 4 if image.size(1) > 5 else self.res # View 1 or View 2? 
+        
+        for i, idx in enumerate(indice):
+            img = image[[i]]
+            img = self.random_crop(idx, img)
+            img = F.interpolate(img, res_tar, mode='bilinear', align_corners=False)
+
+            new_image.append(img)
+
+        new_image = torch.cat(new_image)
+        
+        return new_image
+
+
+
+            
+            
+            
+            
+
+
+
+
+    

libs/data_coco_stuff.py

@@ -0,0 +1,166 @@
+import cv2
+import torch
+from PIL import Image
+import os.path as osp
+import numpy as np
+from torch.utils import data
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+import random 
+
+class RandomResizedCrop(object):
+    def __init__(self, N, res, scale=(0.5, 1.0)):
+        self.res    = res
+        self.scale  = scale 
+        self.rscale = [np.random.uniform(*scale) for _ in range(N)]
+        self.rcrop  = [(np.random.uniform(0, 1), np.random.uniform(0, 1)) for _ in range(N)]
+
+    def random_crop(self, idx, img):
+        ws, hs = self.rcrop[idx]
+        res1 = int(img.size(-1))
+        res2 = int(self.rscale[idx]*res1)
+        i1 = int(round((res1-res2)*ws))
+        j1 = int(round((res1-res2)*hs))
+
+        return img[:, :, i1:i1+res2, j1:j1+res2]
+
+
+    def __call__(self, indice, image):
+        new_image = []
+        res_tar   = self.res // 4 if image.size(1) > 5 else self.res # View 1 or View 2? 
+        
+        for i, idx in enumerate(indice):
+            img = image[[i]]
+            img = self.random_crop(idx, img)
+            img = F.interpolate(img, res_tar, mode='bilinear', align_corners=False)
+
+            new_image.append(img)
+
+        new_image = torch.cat(new_image)
+        
+        return new_image
+
+class RandomVerticalFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        
+    def __call__(self, indice, image):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+
+        if len(image.size()) == 3:
+            image_t = image[I].flip([1]) 
+        else:
+            image_t = image[I].flip([2])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class RandomHorizontalTensorFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, indice, image, is_label=False):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+        
+        if len(image.size()) == 3:
+            image_t = image[I].flip([2])
+        else:
+            image_t = image[I].flip([3])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class _Coco164kCuratedFew(data.Dataset):
+    """Base class
+    This contains fields and methods common to all COCO 164k curated few datasets:
+    
+    (curated) Coco164kFew_Stuff
+    (curated) Coco164kFew_Stuff_People
+    (curated) Coco164kFew_Stuff_Animals
+    (curated) Coco164kFew_Stuff_People_Animals 
+    
+    """
+    def __init__(self, root, img_size, crop_size, split = "train2017"):
+        super(_Coco164kCuratedFew, self).__init__()
+        
+        # work out name
+        self.split = split
+        self.root = root
+        self.include_things_labels = False  # people
+        self.incl_animal_things = False  # animals
+        
+        version = 6
+        
+        name = "Coco164kFew_Stuff"
+        if self.include_things_labels and self.incl_animal_things:
+          name += "_People_Animals"
+        elif self.include_things_labels:
+          name += "_People"
+        elif self.incl_animal_things:
+          name += "_Animals"
+        
+        self.name = (name + "_%d" % version)
+        
+        print("Specific type of _Coco164kCuratedFew dataset: %s" % self.name)
+        
+        self._set_files()
+        
+        
+        self.transform = transforms.Compose([
+                         transforms.RandomChoice([
+                            transforms.ColorJitter(brightness=0.05),
+                            transforms.ColorJitter(contrast=0.05),
+                            transforms.ColorJitter(saturation=0.01),
+                            transforms.ColorJitter(hue=0.01)]),
+                         transforms.RandomHorizontalFlip(),
+                         transforms.RandomVerticalFlip(),
+                         transforms.Resize(int(img_size)),
+                         transforms.RandomCrop(crop_size)])
+
+        N = len(self.files)
+        self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
+        self.random_vertical_flip   = RandomVerticalFlip(N=N)
+        self.random_resized_crop    = RandomResizedCrop(N=N, res=self.res1, scale=self.scale)
+    
+
+    def _set_files(self):
+        # Create data list by parsing the "images" folder
+        if self.split in ["train2017", "val2017"]:
+            file_list = osp.join(self.root, "curated", self.split, self.name + ".txt")
+            file_list = tuple(open(file_list, "r"))
+            file_list = [id_.rstrip() for id_ in file_list]
+            
+            self.files = file_list
+            print("In total {} images.".format(len(self.files)))
+        else:
+            raise ValueError("Invalid split name: {}".format(self.split))
+
+    def __getitem__(self, index):
+        # same as _Coco164k
+        # Set paths
+        image_id = self.files[index]
+        image_path = osp.join(self.root, "images", self.split, image_id + ".jpg")
+        label_path = osp.join(self.root, "annotations", self.split,
+                              image_id + ".png")
+        # Load an image
+        #image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(np.uint8)
+        ori_img = Image.open(image_path)
+        ori_img = self.transform(ori_img)
+        ori_img = np.array(ori_img)
+        if ori_img.ndim < 3:
+            ori_img = np.expand_dims(ori_img, axis=2).repeat(3, axis = 2)
+        ori_img = ori_img[:, :, :3]
+        ori_img = torch.from_numpy(ori_img).float().permute(2, 0, 1)
+        ori_img = ori_img / 255.0
+
+        #label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE).astype(np.int32)
+
+        #label[label == 255] = -1  # to be consistent with 10k
+
+        rets = []
+        rets.append(ori_img)
+        #rets.append(label)
+        return rets
+
+    def __len__(self):
+        return len(self.files)

libs/data_coco_stuff_geo_pho.py

@@ -0,0 +1,145 @@
+import cv2
+import torch
+from PIL import Image
+import os.path as osp
+import numpy as np
+from torch.utils import data
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+import torchvision.transforms.functional as TF
+from .custom_transform import *
+
+class _Coco164kCuratedFew(data.Dataset):
+    """Base class
+    This contains fields and methods common to all COCO 164k curated few datasets:
+    
+    (curated) Coco164kFew_Stuff
+    (curated) Coco164kFew_Stuff_People
+    (curated) Coco164kFew_Stuff_Animals
+    (curated) Coco164kFew_Stuff_People_Animals 
+    
+    """
+    def __init__(self, root, img_size, crop_size, split = "train2017"):
+        super(_Coco164kCuratedFew, self).__init__()
+        
+        # work out name
+        self.split = split
+        self.root = root
+        self.include_things_labels = False  # people
+        self.incl_animal_things = False  # animals
+        
+        version = 6
+        
+        name = "Coco164kFew_Stuff"
+        if self.include_things_labels and self.incl_animal_things:
+          name += "_People_Animals"
+        elif self.include_things_labels:
+          name += "_People"
+        elif self.incl_animal_things:
+          name += "_Animals"
+        
+        self.name = (name + "_%d" % version)
+        
+        print("Specific type of _Coco164kCuratedFew dataset: %s" % self.name)
+        
+        self._set_files()
+        
+        self.transform = transforms.Compose([
+                         transforms.Resize(int(img_size)),
+                         transforms.RandomCrop(crop_size)])
+
+        N = len(self.files)
+        # eqv transform
+        self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
+        self.random_vertical_flip   = RandomVerticalFlip(N=N)
+        self.random_resized_crop    = RandomResizedCrop(N=N, res=288)
+
+        # photometric transform
+        self.random_color_brightness = [RandomColorBrightness(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)]
+        self.random_color_contrast   = [RandomColorContrast(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
+        self.random_color_saturation = [RandomColorSaturation(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
+        self.random_color_hue        = [RandomColorHue(x=0.1, p=0.8, N=N) for _ in range(2)]      # Control this later (NOTE)
+        self.random_gray_scale    = [RandomGrayScale(p=0.2, N=N) for _ in range(2)]
+        self.random_gaussian_blur = [RandomGaussianBlur(sigma=[.1, 2.], p=0.5, N=N) for _ in range(2)]
+
+        self.eqv_list = ['random_crop', 'h_flip']
+        self.inv_list = ['brightness', 'contrast', 'saturation', 'hue', 'gray', 'blur']
+
+        self.transform_tensor = TensorTransform()
+    
+
+    def _set_files(self):
+        # Create data list by parsing the "images" folder
+        if self.split in ["train2017", "val2017"]:
+            file_list = osp.join(self.root, "curated", self.split, self.name + ".txt")
+            file_list = tuple(open(file_list, "r"))
+            file_list = [id_.rstrip() for id_ in file_list]
+            
+            self.files = file_list
+            print("In total {} images.".format(len(self.files)))
+        else:
+            raise ValueError("Invalid split name: {}".format(self.split))
+
+    def transform_eqv(self, indice, image):
+        if 'random_crop' in self.eqv_list:
+            image = self.random_resized_crop(indice, image)
+        if 'h_flip' in self.eqv_list:
+            image = self.random_horizontal_flip(indice, image)
+        if 'v_flip' in self.eqv_list:
+            image = self.random_vertical_flip(indice, image)
+
+        return image
+    
+    def transform_inv(self, index, image, ver):
+        """
+        Hyperparameters same as MoCo v2. 
+        (https://github.com/facebookresearch/moco/blob/master/main_moco.py)
+        """
+        if 'brightness' in self.inv_list:
+            image = self.random_color_brightness[ver](index, image)
+        if 'contrast' in self.inv_list:
+            image = self.random_color_contrast[ver](index, image)
+        if 'saturation' in self.inv_list:
+            image = self.random_color_saturation[ver](index, image)
+        if 'hue' in self.inv_list:
+            image = self.random_color_hue[ver](index, image)
+        if 'gray' in self.inv_list:
+            image = self.random_gray_scale[ver](index, image)
+        if 'blur' in self.inv_list:
+            image = self.random_gaussian_blur[ver](index, image)
+        
+        return image
+    
+    def transform_image(self, index, image):
+        image1 = self.transform_inv(index, image, 0)
+        image1 = self.transform_tensor(image)
+
+        image2 = self.transform_inv(index, image, 1)
+        #image2 = TF.resize(image2, self.crop_size, Image.BILINEAR)
+        image2 = self.transform_tensor(image2)
+        return image1, image2
+
+    def __getitem__(self, index):
+        # same as _Coco164k
+        # Set paths
+        image_id = self.files[index]
+        image_path = osp.join(self.root, "images", self.split, image_id + ".jpg")
+        # Load an image
+        ori_img = Image.open(image_path)
+        ori_img = self.transform(ori_img)
+
+        image1, image2 = self.transform_image(index, ori_img)
+        if image1.shape[0] < 3:
+            image1 = image1.repeat(3, 1, 1)
+        if image2.shape[0] < 3:
+            image2 = image2.repeat(3, 1, 1)
+
+        rets = []
+        rets.append(image1)
+        rets.append(image2)
+        rets.append(index)
+        
+        return rets
+
+    def __len__(self):
+        return len(self.files)

libs/data_geo.py

@@ -0,0 +1,176 @@
+"""SLIC dataset
+ - Returns an image together with its SLIC segmentation map.
+"""
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+import numpy as np
+from glob import glob
+from PIL import Image
+from skimage.segmentation import slic
+from skimage.color import rgb2lab
+import torch.nn.functional as F
+
+from .utils import label2one_hot_torch
+
+class RandomResizedCrop(object):
+    def __init__(self, N, res, scale=(0.5, 1.0)):
+        self.res    = res
+        self.scale  = scale 
+        self.rscale = [np.random.uniform(*scale) for _ in range(N)]
+        self.rcrop  = [(np.random.uniform(0, 1), np.random.uniform(0, 1)) for _ in range(N)]
+
+    def random_crop(self, idx, img):
+        ws, hs = self.rcrop[idx]
+        res1 = int(img.size(-1))
+        res2 = int(self.rscale[idx]*res1)
+        i1 = int(round((res1-res2)*ws))
+        j1 = int(round((res1-res2)*hs))
+
+        return img[:, :, i1:i1+res2, j1:j1+res2]
+
+
+    def __call__(self, indice, image):
+        new_image = []
+        res_tar   = self.res // 8 if image.size(1) > 5 else self.res # View 1 or View 2? 
+        
+        for i, idx in enumerate(indice):
+            img = image[[i]]
+            img = self.random_crop(idx, img)
+            img = F.interpolate(img, res_tar, mode='bilinear', align_corners=False)
+
+            new_image.append(img)
+
+        new_image = torch.cat(new_image)
+        
+        return new_image
+
+class RandomVerticalFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        
+    def __call__(self, indice, image):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+
+        if len(image.size()) == 3:
+            image_t = image[I].flip([1]) 
+        else:
+            image_t = image[I].flip([2])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class RandomHorizontalTensorFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, indice, image, is_label=False):
+        I = np.nonzero(self.plist[indice.cpu()] < self.p_ref)[0]
+        
+        if len(image.size()) == 3:
+            image_t = image[I].flip([2])
+        else:
+            image_t = image[I].flip([3])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class Dataset(data.Dataset):
+    def __init__(self, data_dir, img_size=256, crop_size=128, test=False, 
+                 sp_num=256, slic = True, lab = False): 
+        super(Dataset, self).__init__()
+        #self.data_list = glob(os.path.join(data_dir, "*.jpg"))
+        ext = ["*.jpg"]
+        dl = []
+        [dl.extend(glob(data_dir + '/**/' + e, recursive=True)) for e in ext]
+        self.data_list = dl
+        self.sp_num = sp_num
+        self.slic = slic
+        self.lab = lab
+        if test:
+            self.transform = transforms.Compose([
+                             transforms.Resize(img_size),
+                             transforms.CenterCrop(crop_size)])
+        else:
+            self.transform = transforms.Compose([
+                             transforms.RandomChoice([
+                                transforms.ColorJitter(brightness=0.05),
+                                transforms.ColorJitter(contrast=0.05),
+                                transforms.ColorJitter(saturation=0.01),
+                                transforms.ColorJitter(hue=0.01)]),
+                             transforms.RandomHorizontalFlip(),
+                             transforms.RandomVerticalFlip(),
+                             transforms.Resize(int(img_size)),
+                             transforms.RandomCrop(crop_size)])
+
+        N = len(self.data_list)
+        self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
+        self.random_vertical_flip   = RandomVerticalFlip(N=N)
+        self.random_resized_crop    = RandomResizedCrop(N=N, res=224)
+        self.eqv_list = ['random_crop', 'h_flip']
+
+    def transform_eqv(self, indice, image):
+        if 'random_crop' in self.eqv_list:
+            image = self.random_resized_crop(indice, image)
+        if 'h_flip' in self.eqv_list:
+            image = self.random_horizontal_flip(indice, image)
+        if 'v_flip' in self.eqv_list:
+            image = self.random_vertical_flip(indice, image)
+
+        return image
+    
+    def __getitem__(self, index):
+        data_path = self.data_list[index]
+        ori_img = Image.open(data_path)
+        ori_img = self.transform(ori_img)
+        ori_img = np.array(ori_img)
+
+        # compute slic
+        if self.slic:
+            slic_i = slic(ori_img, n_segments=self.sp_num, compactness=10, start_label=0, min_size_factor=0.3)
+            slic_i = torch.from_numpy(slic_i)
+            slic_i[slic_i >= self.sp_num] = self.sp_num - 1
+            oh = label2one_hot_torch(slic_i.unsqueeze(0).unsqueeze(0), C = self.sp_num).squeeze()
+
+        if ori_img.ndim < 3:
+            ori_img = np.expand_dims(ori_img, axis=2).repeat(3, axis = 2)
+        ori_img = ori_img[:, :, :3]
+
+        rets = []
+        if self.lab:
+            lab_img = rgb2lab(ori_img)
+            rets.append(torch.from_numpy(lab_img).float().permute(2, 0, 1))
+
+        ori_img = torch.from_numpy(ori_img).float().permute(2, 0, 1)
+        rets.append(ori_img/255.0)
+
+        if self.slic:
+            rets.append(oh)
+        
+        rets.append(index)
+        
+        return rets
+    
+    def __len__(self):
+        return len(self.data_list)
+
+if __name__ == '__main__':
+    import torchvision.utils as vutils
+    dataset = Dataset('/home/xtli/DATA/texture_data/',
+                      sampled_num=3000)
+    loader_ = torch.utils.data.DataLoader(dataset     = dataset,
+                                         batch_size  = 1,
+                                         shuffle     = True,
+                                         num_workers = 1,
+                                         drop_last   = True)
+    loader = iter(loader_)
+    img, points, pixs = loader.next()
+
+    crop_size = 128
+    canvas = torch.zeros((1, 3, crop_size, crop_size))
+    for i in range(points.shape[-2]):
+        p = (points[0, i] + 1) / 2.0 * (crop_size - 1)
+        canvas[0, :, int(p[0]), int(p[1])] = pixs[0, :, i]
+    vutils.save_image(canvas, 'canvas.png')
+    vutils.save_image(img, 'img.png')

libs/data_geo_pho.py

@@ -0,0 +1,130 @@
+"""SLIC dataset
+ - Returns an image together with its SLIC segmentation map.
+"""
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+import numpy as np
+from glob import glob
+from PIL import Image
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+
+from .custom_transform import *
+
+class Dataset(data.Dataset):
+    def __init__(self, data_dir, img_size=256, crop_size=128, test=False, 
+                 sp_num=256, slic = True, lab = False): 
+        super(Dataset, self).__init__()
+        #self.data_list = glob(os.path.join(data_dir, "*.jpg"))
+        ext = ["*.jpg"]
+        dl = []
+        [dl.extend(glob(data_dir + '/**/' + e, recursive=True)) for e in ext]
+        self.data_list = dl
+        self.sp_num = sp_num
+        self.slic = slic
+        self.lab = lab
+        if test:
+            self.transform = transforms.Compose([
+                             transforms.Resize(img_size),
+                             transforms.CenterCrop(crop_size)])
+        else:
+            self.transform = transforms.Compose([
+                             transforms.Resize(int(img_size)),
+                             transforms.RandomCrop(crop_size)])
+
+        N = len(self.data_list)
+        # eqv transform
+        self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
+        self.random_vertical_flip   = RandomVerticalFlip(N=N)
+        self.random_resized_crop    = RandomResizedCrop(N=N, res=256)
+
+        # photometric transform
+        self.random_color_brightness = [RandomColorBrightness(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)]
+        self.random_color_contrast   = [RandomColorContrast(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
+        self.random_color_saturation = [RandomColorSaturation(x=0.3, p=0.8, N=N) for _ in range(2)] # Control this later (NOTE)
+        self.random_color_hue        = [RandomColorHue(x=0.1, p=0.8, N=N) for _ in range(2)]      # Control this later (NOTE)
+        self.random_gray_scale    = [RandomGrayScale(p=0.2, N=N) for _ in range(2)]
+        self.random_gaussian_blur = [RandomGaussianBlur(sigma=[.1, 2.], p=0.5, N=N) for _ in range(2)]
+
+        self.eqv_list = ['random_crop', 'h_flip']
+        self.inv_list = ['brightness', 'contrast', 'saturation', 'hue', 'gray', 'blur']
+
+        self.transform_tensor = TensorTransform()
+
+    def transform_eqv(self, indice, image):
+        if 'random_crop' in self.eqv_list:
+            image = self.random_resized_crop(indice, image)
+        if 'h_flip' in self.eqv_list:
+            image = self.random_horizontal_flip(indice, image)
+        if 'v_flip' in self.eqv_list:
+            image = self.random_vertical_flip(indice, image)
+
+        return image
+    
+    def transform_inv(self, index, image, ver):
+        """
+        Hyperparameters same as MoCo v2. 
+        (https://github.com/facebookresearch/moco/blob/master/main_moco.py)
+        """
+        if 'brightness' in self.inv_list:
+            image = self.random_color_brightness[ver](index, image)
+        if 'contrast' in self.inv_list:
+            image = self.random_color_contrast[ver](index, image)
+        if 'saturation' in self.inv_list:
+            image = self.random_color_saturation[ver](index, image)
+        if 'hue' in self.inv_list:
+            image = self.random_color_hue[ver](index, image)
+        if 'gray' in self.inv_list:
+            image = self.random_gray_scale[ver](index, image)
+        if 'blur' in self.inv_list:
+            image = self.random_gaussian_blur[ver](index, image)
+        
+        return image
+    
+    def transform_image(self, index, image):
+        image1 = self.transform_inv(index, image, 0)
+        image1 = self.transform_tensor(image)
+
+        image2 = self.transform_inv(index, image, 1)
+        #image2 = TF.resize(image2, self.crop_size, Image.BILINEAR)
+        image2 = self.transform_tensor(image2)
+        return image1, image2
+    
+    def __getitem__(self, index):
+        data_path = self.data_list[index]
+        ori_img = Image.open(data_path)
+        ori_img = self.transform(ori_img)
+
+        image1, image2 = self.transform_image(index, ori_img)
+
+        rets = []
+        rets.append(image1)
+        rets.append(image2)
+        rets.append(index)
+        
+        return rets
+    
+    def __len__(self):
+        return len(self.data_list)
+
+if __name__ == '__main__':
+    import torchvision.utils as vutils
+    dataset = Dataset('/home/xtli/DATA/texture_data/',
+                      sampled_num=3000)
+    loader_ = torch.utils.data.DataLoader(dataset     = dataset,
+                                         batch_size  = 1,
+                                         shuffle     = True,
+                                         num_workers = 1,
+                                         drop_last   = True)
+    loader = iter(loader_)
+    img, points, pixs = loader.next()
+
+    crop_size = 128
+    canvas = torch.zeros((1, 3, crop_size, crop_size))
+    for i in range(points.shape[-2]):
+        p = (points[0, i] + 1) / 2.0 * (crop_size - 1)
+        canvas[0, :, int(p[0]), int(p[1])] = pixs[0, :, i]
+    vutils.save_image(canvas, 'canvas.png')
+    vutils.save_image(img, 'img.png')

libs/data_slic.py

@@ -0,0 +1,175 @@
+"""SLIC dataset
+ - Returns an image together with its SLIC segmentation map.
+"""
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+import numpy as np
+from glob import glob
+from PIL import Image
+from skimage.segmentation import slic
+from skimage.color import rgb2lab
+
+from .utils import label2one_hot_torch
+
+class RandomResizedCrop(object):
+    def __init__(self, N, res, scale=(0.5, 1.0)):
+        self.res    = res
+        self.scale  = scale 
+        self.rscale = [np.random.uniform(*scale) for _ in range(N)]
+        self.rcrop  = [(np.random.uniform(0, 1), np.random.uniform(0, 1)) for _ in range(N)]
+
+    def random_crop(self, idx, img):
+        ws, hs = self.rcrop[idx]
+        res1 = int(img.size(-1))
+        res2 = int(self.rscale[idx]*res1)
+        i1 = int(round((res1-res2)*ws))
+        j1 = int(round((res1-res2)*hs))
+
+        return img[:, :, i1:i1+res2, j1:j1+res2]
+
+
+    def __call__(self, indice, image):
+        new_image = []
+        res_tar   = self.res // 4 if image.size(1) > 5 else self.res # View 1 or View 2? 
+        
+        for i, idx in enumerate(indice):
+            img = image[[i]]
+            img = self.random_crop(idx, img)
+            img = F.interpolate(img, res_tar, mode='bilinear', align_corners=False)
+
+            new_image.append(img)
+
+        new_image = torch.cat(new_image)
+        
+        return new_image
+
+class RandomVerticalFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+        
+    def __call__(self, indice, image):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+
+        if len(image.size()) == 3:
+            image_t = image[I].flip([1]) 
+        else:
+            image_t = image[I].flip([2])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class RandomHorizontalTensorFlip(object):
+    def __init__(self, N, p=0.5):
+        self.p_ref = p
+        self.plist = np.random.random_sample(N)
+
+    def __call__(self, indice, image, is_label=False):
+        I = np.nonzero(self.plist[indice] < self.p_ref)[0]
+        
+        if len(image.size()) == 3:
+            image_t = image[I].flip([2])
+        else:
+            image_t = image[I].flip([3])
+        
+        return torch.stack([image_t[np.where(I==i)[0][0]] if i in I else image[i] for i in range(image.size(0))])
+
+class Dataset(data.Dataset):
+    def __init__(self, data_dir, img_size=256, crop_size=128, test=False, 
+                 sp_num=256, slic = True, lab = False): 
+        super(Dataset, self).__init__()
+        #self.data_list = glob(os.path.join(data_dir, "*.jpg"))
+        ext = ["*.jpg"]
+        dl = []
+        [dl.extend(glob(data_dir + '/**/' + e, recursive=True)) for e in ext]
+        self.data_list = dl
+        self.sp_num = sp_num
+        self.slic = slic
+        self.lab = lab
+        if test:
+            self.transform = transforms.Compose([
+                             transforms.Resize(img_size),
+                             transforms.CenterCrop(crop_size)])
+        else:
+            self.transform = transforms.Compose([
+                             transforms.RandomChoice([
+                                transforms.ColorJitter(brightness=0.05),
+                                transforms.ColorJitter(contrast=0.05),
+                                transforms.ColorJitter(saturation=0.01),
+                                transforms.ColorJitter(hue=0.01)]),
+                             transforms.RandomHorizontalFlip(),
+                             transforms.RandomVerticalFlip(),
+                             transforms.Resize(int(img_size)),
+                             transforms.RandomCrop(crop_size)])
+
+        N = len(self.data_list)
+        self.random_horizontal_flip = RandomHorizontalTensorFlip(N=N)
+        self.random_vertical_flip   = RandomVerticalFlip(N=N)
+        self.random_resized_crop    = RandomResizedCrop(N=N, res=img_size)
+        self.eqv_list = ['random_crop', 'h_flip']
+
+    def transform_eqv(self, indice, image):
+        if 'random_crop' in self.eqv_list:
+            image = self.random_resized_crop(indice, image)
+        if 'h_flip' in self.eqv_list:
+            image = self.random_horizontal_flip(indice, image)
+        if 'v_flip' in self.eqv_list:
+            image = self.random_vertical_flip(indice, image)
+
+        return image
+    
+    def __getitem__(self, index):
+        data_path = self.data_list[index]
+        ori_img = Image.open(data_path)
+        ori_img = self.transform(ori_img)
+        ori_img = np.array(ori_img)
+
+        # compute slic
+        if self.slic:
+            slic_i = slic(ori_img, n_segments=self.sp_num, compactness=10, start_label=0, min_size_factor=0.3)
+            slic_i = torch.from_numpy(slic_i)
+            slic_i[slic_i >= self.sp_num] = self.sp_num - 1
+            oh = label2one_hot_torch(slic_i.unsqueeze(0).unsqueeze(0), C = self.sp_num).squeeze()
+
+        if ori_img.ndim < 3:
+            ori_img = np.expand_dims(ori_img, axis=2).repeat(3, axis = 2)
+        ori_img = ori_img[:, :, :3]
+
+        rets = []
+        if self.lab:
+            lab_img = rgb2lab(ori_img)
+            rets.append(torch.from_numpy(lab_img).float().permute(2, 0, 1))
+
+        ori_img = torch.from_numpy(ori_img).float().permute(2, 0, 1)
+        rets.append(ori_img/255.0)
+
+        if self.slic:
+            rets.append(oh)
+        
+        rets.append(index)
+        
+        return rets
+    
+    def __len__(self):
+        return len(self.data_list)
+
+if __name__ == '__main__':
+    import torchvision.utils as vutils
+    dataset = Dataset('/home/xtli/DATA/texture_data/',
+                      sampled_num=3000)
+    loader_ = torch.utils.data.DataLoader(dataset     = dataset,
+                                         batch_size  = 1,
+                                         shuffle     = True,
+                                         num_workers = 1,
+                                         drop_last   = True)
+    loader = iter(loader_)
+    img, points, pixs = loader.next()
+
+    crop_size = 128
+    canvas = torch.zeros((1, 3, crop_size, crop_size))
+    for i in range(points.shape[-2]):
+        p = (points[0, i] + 1) / 2.0 * (crop_size - 1)
+        canvas[0, :, int(p[0]), int(p[1])] = pixs[0, :, i]
+    vutils.save_image(canvas, 'canvas.png')
+    vutils.save_image(img, 'img.png')

libs/discriminator.py

@@ -0,0 +1,60 @@
+import functools
+import torch.nn as nn
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        norm_layer = nn.BatchNorm2d
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)

libs/flow_transforms.py

@@ -0,0 +1,393 @@
+from __future__ import division
+import torch
+import random
+import numpy as np
+import numbers
+import types
+import scipy.ndimage as ndimage
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image
+# import torchvision.transforms.functional as FF
+
+'''
+Data argumentation file 
+modifed from 
+https://github.com/ClementPinard/FlowNetPytorch
+
+
+'''
+
+
+
+'''Set of tranform random routines that takes both input and target as arguments,
+in order to have random but coherent transformations.
+inputs are PIL Image pairs and targets are ndarrays'''
+
+_pil_interpolation_to_str = {
+    Image.NEAREST: 'PIL.Image.NEAREST',
+    Image.BILINEAR: 'PIL.Image.BILINEAR',
+    Image.BICUBIC: 'PIL.Image.BICUBIC',
+    Image.LANCZOS: 'PIL.Image.LANCZOS',
+    Image.HAMMING: 'PIL.Image.HAMMING',
+    Image.BOX: 'PIL.Image.BOX',
+}
+
+class Compose(object):
+    """ Composes several co_transforms together.
+    For example:
+    >>> co_transforms.Compose([
+    >>>     co_transforms.CenterCrop(10),
+    >>>     co_transforms.ToTensor(),
+    >>>  ])
+    """
+
+    def __init__(self, co_transforms):
+        self.co_transforms = co_transforms
+
+    def __call__(self, input, target):
+        for t in self.co_transforms:
+            input,target = t(input,target)
+        return input,target
+
+
+class ArrayToTensor(object):
+    """Converts a numpy.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W)."""
+
+    def __call__(self, array):
+        assert(isinstance(array, np.ndarray))
+
+        array = np.transpose(array, (2, 0, 1))
+        # handle numpy array
+        tensor = torch.from_numpy(array)
+        # put it from HWC to CHW format
+
+        return tensor.float()
+
+
+class ArrayToPILImage(object):
+    """Converts a numpy.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W)."""
+
+    def __call__(self, array):
+        assert(isinstance(array, np.ndarray))
+
+        img = Image.fromarray(array.astype(np.uint8))
+
+        return img
+
+class PILImageToTensor(object):
+    """Converts a numpy.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W)."""
+
+    def __call__(self, img):
+        assert(isinstance(img, Image.Image))
+
+        array = np.asarray(img)
+        array = np.transpose(array, (2, 0, 1))
+        tensor = torch.from_numpy(array)
+
+        return tensor.float()
+
+
+class Lambda(object):
+    """Applies a lambda as a transform"""
+
+    def __init__(self, lambd):
+        assert isinstance(lambd, types.LambdaType)
+        self.lambd = lambd
+
+    def __call__(self, input,target):
+        return self.lambd(input,target)
+
+
+class CenterCrop(object):
+    """Crops the given inputs and target arrays at the center to have a region of
+    the given size. size can be a tuple (target_height, target_width)
+    or an integer, in which case the target will be of a square shape (size, size)
+    Careful, img1 and img2 may not be the same size
+    """
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, inputs, target):
+        h1, w1, _ = inputs[0].shape
+        # h2, w2, _ = inputs[1].shape
+        th, tw = self.size
+        x1 = int(round((w1 - tw) / 2.))
+        y1 = int(round((h1 - th) / 2.))
+        # x2 = int(round((w2 - tw) / 2.))
+        # y2 = int(round((h2 - th) / 2.))
+        for i in range(len(inputs)):
+            inputs[i] = inputs[i][y1: y1 + th, x1: x1 + tw]
+        # inputs[0] = inputs[0][y1: y1 + th, x1: x1 + tw]
+        # inputs[1] = inputs[1][y2: y2 + th, x2: x2 + tw]
+        target = target[y1: y1 + th, x1: x1 + tw]
+        return inputs,target
+
+class myRandomResized(object):
+    """
+    based on RandomResizedCrop in
+    https://pytorch.org/docs/stable/_modules/torchvision/transforms/transforms.html#RandomResizedCrop
+    """
+
+    def __init__(self, expect_min_size, scale=(0.8, 1.5), interpolation=cv2.INTER_NEAREST):
+        # assert (min(input_size) * min(scale) > max(expect_size))
+        # one consider one decimal !!
+        assert (isinstance(scale,tuple) and len(scale)==2)
+        self.interpolation = interpolation
+        self.scale = [ x*0.1 for x in range(int(scale[0]*10),int(scale[1])*10 )]
+        self.min_size = expect_min_size
+
+    @staticmethod
+    def get_params(img, scale, min_size):
+        """Get parameters for ``crop`` for a random sized crop.
+
+        Args:
+            img (PIL Image): Image to be cropped.
+            scale (tuple): range of size of the origin size cropped
+            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped
+
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
+                sized crop.
+        """
+        # area = img.size[0] * img.size[1]
+        h, w, _ = img.shape
+        for attempt in range(10):
+            rand_scale_ = random.choice(scale)
+
+            if random.random() < 0.5:
+                rand_scale = rand_scale_
+            else:
+                rand_scale = -1.
+
+            if min_size[0] <= rand_scale * h and min_size[1] <= rand_scale * w\
+                    and rand_scale * h % 16 == 0 and rand_scale * w %16 ==0 :
+                # the 16*n condition is for network architecture
+                return (int(rand_scale * h),int(rand_scale * w ))
+
+        # Fallback
+        return (h, w)
+
+    def __call__(self, inputs, tgt):
+        """
+        Args:
+            img (PIL Image): Image to be cropped and resized.
+
+        Returns:
+            PIL Image: Randomly cropped and resized image.
+        """
+        h,w = self.get_params(inputs[0], self.scale, self.min_size)
+        for i in range(len(inputs)):
+            inputs[i] =  cv2.resize(inputs[i], (w,h), self.interpolation)
+
+        tgt =  cv2.resize(tgt, (w,h), self.interpolation) #for input as h*w*1 the output is h*w
+        return inputs, np.expand_dims(tgt,-1)
+
+    def __repr__(self):
+        interpolate_str = _pil_interpolation_to_str[self.interpolation]
+        format_string = self.__class__.__name__ + '(min_size={0}'.format(self.min_size)
+        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
+        format_string += ', interpolation={0})'.format(interpolate_str)
+        return format_string
+
+
+class Scale(object):
+    """ Rescales the inputs and target arrays to the given 'size'.
+    'size' will be the size of the smaller edge.
+    For example, if height > width, then image will be
+    rescaled to (size * height / width, size)
+    size: size of the smaller edge
+    interpolation order: Default: 2 (bilinear)
+    """
+
+    def __init__(self, size, order=2):
+        self.size = size
+        self.order = order
+
+    def __call__(self, inputs, target):
+        h, w, _ = inputs[0].shape
+        if (w <= h and w == self.size) or (h <= w and h == self.size):
+            return inputs,target
+        if w < h:
+            ratio = self.size/w
+        else:
+            ratio = self.size/h
+
+        for i in range(len(inputs)):
+            inputs[i] = ndimage.interpolation.zoom(inputs[i], ratio, order=self.order)[:, :, :3]
+
+        target = ndimage.interpolation.zoom(target, ratio, order=self.order)[:, :, :1]
+        #target *= ratio
+        return inputs, target
+
+
+class RandomCrop(object):
+    """Crops the given PIL.Image at a random location to have a region of
+    the given size. size can be a tuple (target_height, target_width)
+    or an integer, in which case the target will be of a square shape (size, size)
+    """
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, inputs,target):
+        h, w, _ = inputs[0].shape
+        th, tw = self.size
+        if w == tw and h == th:
+            return inputs,target
+
+        x1 = random.randint(0, w - tw)
+        y1 = random.randint(0, h - th)
+        for i in range(len(inputs)):
+            inputs[i] = inputs[i][y1: y1 + th,x1: x1 + tw]
+            # inputs[1] = inputs[1][y1: y1 + th,x1: x1 + tw]
+            # inputs[2] = inputs[2][y1: y1 + th, x1: x1 + tw]
+
+        return inputs, target[y1: y1 + th,x1: x1 + tw]
+
+class MyScale(object):
+    def __init__(self, size, order=2):
+        self.size = size
+        self.order = order
+
+    def __call__(self, inputs, target):
+        h, w, _ = inputs[0].shape
+        if (w <= h and w == self.size) or (h <= w and h == self.size):
+            return inputs,target
+        if w < h:
+            for i in range(len(inputs)):
+                inputs[i] = cv2.resize(inputs[i], (self.size, int(h * self.size / w)))
+            target = cv2.resize(target.squeeze(), (self.size, int(h * self.size / w)), cv2.INTER_NEAREST)
+        else:
+            for i in range(len(inputs)):
+                inputs[i] = cv2.resize(inputs[i], (int(w * self.size / h), self.size))
+            target = cv2.resize(target.squeeze(), (int(w * self.size / h), self.size), cv2.INTER_NEAREST)
+        target = np.expand_dims(target, axis=2)
+        return inputs, target
+
+class RandomHorizontalFlip(object):
+    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
+    """
+
+    def __call__(self, inputs, target):
+        if random.random() < 0.5:
+            for i in range(len(inputs)):
+                inputs[i] = np.copy(np.fliplr(inputs[i]))
+                # inputs[1] = np.copy(np.fliplr(inputs[1]))
+                # inputs[2] = np.copy(np.fliplr(inputs[2]))
+
+            target = np.copy(np.fliplr(target))
+            # target[:,:,0] *= -1
+        return inputs,target
+
+
+class RandomVerticalFlip(object):
+    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
+    """
+
+    def __call__(self, inputs, target):
+        if random.random() < 0.5:
+            for i in range(len(inputs)):
+                inputs[i] = np.copy(np.flipud(inputs[i]))
+                # inputs[1] = np.copy(np.flipud(inputs[1]))
+                # inputs[2] = np.copy(np.flipud(inputs[2]))
+
+            target = np.copy(np.flipud(target))
+            # target[:,:,1] *= -1 #for disp there is no y dim
+        return inputs,target
+
+
+class RandomRotate(object):
+    """Random rotation of the image from -angle to angle (in degrees)
+    This is useful for dataAugmentation, especially for geometric problems such as FlowEstimation
+    angle: max angle of the rotation
+    interpolation order: Default: 2 (bilinear)
+    reshape: Default: false. If set to true, image size will be set to keep every pixel in the image.
+    diff_angle: Default: 0. Must stay less than 10 degrees, or linear approximation of flowmap will be off.
+    """
+
+    def __init__(self, angle, diff_angle=0, order=2, reshape=False):
+        self.angle = angle
+        self.reshape = reshape
+        self.order = order
+        self.diff_angle = diff_angle
+
+    def __call__(self, inputs,target):
+        applied_angle = random.uniform(-self.angle,self.angle)
+        diff = random.uniform(-self.diff_angle,self.diff_angle)
+        angle1 = applied_angle - diff/2
+        angle2 = applied_angle + diff/2
+        angle1_rad = angle1*np.pi/180
+
+        h, w, _ = target.shape
+
+        def rotate_flow(i,j,k):
+            return -k*(j-w/2)*(diff*np.pi/180) + (1-k)*(i-h/2)*(diff*np.pi/180)
+
+        rotate_flow_map = np.fromfunction(rotate_flow, target.shape)
+        target += rotate_flow_map
+
+        inputs[0] = ndimage.interpolation.rotate(inputs[0], angle1, reshape=self.reshape, order=self.order)
+        inputs[1] = ndimage.interpolation.rotate(inputs[1], angle2, reshape=self.reshape, order=self.order)
+        target = ndimage.interpolation.rotate(target, angle1, reshape=self.reshape, order=self.order)
+        # flow vectors must be rotated too! careful about Y flow which is upside down
+        target_ = np.copy(target)
+        target[:,:,0] = np.cos(angle1_rad)*target_[:,:,0] + np.sin(angle1_rad)*target_[:,:,1]
+        target[:,:,1] = -np.sin(angle1_rad)*target_[:,:,0] + np.cos(angle1_rad)*target_[:,:,1]
+        return inputs,target
+
+
+class RandomTranslate(object):
+    def __init__(self, translation):
+        if isinstance(translation, numbers.Number):
+            self.translation = (int(translation), int(translation))
+        else:
+            self.translation = translation
+
+    def __call__(self, inputs,target):
+        h, w, _ = inputs[0].shape
+        th, tw = self.translation
+        tw = random.randint(-tw, tw)
+        th = random.randint(-th, th)
+        if tw == 0 and th == 0:
+            return inputs, target
+        # compute x1,x2,y1,y2 for img1 and target, and x3,x4,y3,y4 for img2
+        x1,x2,x3,x4 = max(0,tw), min(w+tw,w), max(0,-tw), min(w-tw,w)
+        y1,y2,y3,y4 = max(0,th), min(h+th,h), max(0,-th), min(h-th,h)
+
+        inputs[0] = inputs[0][y1:y2,x1:x2]
+        inputs[1] = inputs[1][y3:y4,x3:x4]
+        target = target[y1:y2,x1:x2]
+        target[:,:,0] += tw
+        target[:,:,1] += th
+
+        return inputs, target
+
+
+class RandomColorWarp(object):
+    def __init__(self, mean_range=0, std_range=0):
+        self.mean_range = mean_range
+        self.std_range = std_range
+
+    def __call__(self, inputs, target):
+        random_std = np.random.uniform(-self.std_range, self.std_range, 3)
+        random_mean = np.random.uniform(-self.mean_range, self.mean_range, 3)
+        random_order = np.random.permutation(3)
+
+        inputs[0] *= (1 + random_std)
+        inputs[0] += random_mean
+
+        inputs[1] *= (1 + random_std)
+        inputs[1] += random_mean
+
+        inputs[0] = inputs[0][:,:,random_order]
+        inputs[1] = inputs[1][:,:,random_order]
+
+        return inputs, target

libs/losses.py

@@ -0,0 +1,416 @@
+from libs.blocks import encoder5
+import torch
+import torchvision
+import torch.nn as nn
+from torch.nn import init
+import torch.nn.functional as F
+from .normalization import get_nonspade_norm_layer
+from .blocks import encoder5
+
+import os
+import numpy as np
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def print_network(self):
+        if isinstance(self, list):
+            self = self[0]
+        num_params = 0
+        for param in self.parameters():
+            num_params += param.numel()
+        print('Network [%s] was created. Total number of parameters: %.1f million. '
+              'To see the architecture, do print(network).'
+              % (type(self).__name__, num_params / 1000000))
+
+    def init_weights(self, init_type='normal', gain=0.02):
+        def init_func(m):
+            classname = m.__class__.__name__
+            if classname.find('BatchNorm2d') != -1:
+                if hasattr(m, 'weight') and m.weight is not None:
+                    init.normal_(m.weight.data, 1.0, gain)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+            elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+                if init_type == 'normal':
+                    init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == 'xavier':
+                    init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == 'xavier_uniform':
+                    init.xavier_uniform_(m.weight.data, gain=1.0)
+                elif init_type == 'kaiming':
+                    init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                elif init_type == 'orthogonal':
+                    init.orthogonal_(m.weight.data, gain=gain)
+                elif init_type == 'none':  # uses pytorch's default init method
+                    m.reset_parameters()
+                else:
+                    raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+        # propagate to children
+        for m in self.children():
+            if hasattr(m, 'init_weights'):
+                m.init_weights(init_type, gain)
+
+class NLayerDiscriminator(BaseNetwork):
+    def __init__(self):
+        super().__init__()
+
+        kw = 4
+        padw = int(np.ceil((kw - 1.0) / 2))
+        nf = 64
+        n_layers_D = 4
+        input_nc = 3
+
+        norm_layer = get_nonspade_norm_layer('spectralinstance')
+        sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
+                     nn.LeakyReLU(0.2, False)]]
+
+        for n in range(1, n_layers_D):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            stride = 1 if n == n_layers_D - 1 else 2
+            sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw,
+                                               stride=stride, padding=padw)),
+                          nn.LeakyReLU(0.2, False)
+                          ]]
+
+        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
+
+        # We divide the layers into groups to extract intermediate layer outputs
+        for n in range(len(sequence)):
+            self.add_module('model' + str(n), nn.Sequential(*sequence[n]))
+
+    def forward(self, input, get_intermediate_features = True):
+        results = [input]
+        for submodel in self.children():
+            intermediate_output = submodel(results[-1])
+            results.append(intermediate_output)
+
+        if get_intermediate_features:
+            return results[1:]
+        else:
+            return results[-1]
+
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super().__init__()
+        vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        import pdb; pdb.set_trace()
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+class encoder5(nn.Module):
+    def __init__(self):
+        super(encoder5,self).__init__()
+        # vgg
+        # 224 x 224
+        self.conv1 = nn.Conv2d(3,3,1,1,0)
+        self.reflecPad1 = nn.ReflectionPad2d((1,1,1,1))
+        # 226 x 226
+
+        self.conv2 = nn.Conv2d(3,64,3,1,0)
+        self.relu2 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.reflecPad3 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv3 = nn.Conv2d(64,64,3,1,0)
+        self.relu3 = nn.ReLU(inplace=True)
+        # 224 x 224
+
+        self.maxPool = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 112 x 112
+
+        self.reflecPad4 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv4 = nn.Conv2d(64,128,3,1,0)
+        self.relu4 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.reflecPad5 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv5 = nn.Conv2d(128,128,3,1,0)
+        self.relu5 = nn.ReLU(inplace=True)
+        # 112 x 112
+
+        self.maxPool2 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 56 x 56
+
+        self.reflecPad6 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv6 = nn.Conv2d(128,256,3,1,0)
+        self.relu6 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad7 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv7 = nn.Conv2d(256,256,3,1,0)
+        self.relu7 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad8 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv8 = nn.Conv2d(256,256,3,1,0)
+        self.relu8 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.reflecPad9 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv9 = nn.Conv2d(256,256,3,1,0)
+        self.relu9 = nn.ReLU(inplace=True)
+        # 56 x 56
+
+        self.maxPool3 = nn.MaxPool2d(kernel_size=2,stride=2)
+        # 28 x 28
+
+        self.reflecPad10 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv10 = nn.Conv2d(256,512,3,1,0)
+        self.relu10 = nn.ReLU(inplace=True)
+
+        self.reflecPad11 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv11 = nn.Conv2d(512,512,3,1,0)
+        self.relu11 = nn.ReLU(inplace=True)
+
+        self.reflecPad12 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv12 = nn.Conv2d(512,512,3,1,0)
+        self.relu12 = nn.ReLU(inplace=True)
+
+        self.reflecPad13 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv13 = nn.Conv2d(512,512,3,1,0)
+        self.relu13 = nn.ReLU(inplace=True)
+
+        self.maxPool4 = nn.MaxPool2d(kernel_size=2,stride=2)
+        self.reflecPad14 = nn.ReflectionPad2d((1,1,1,1))
+        self.conv14 = nn.Conv2d(512,512,3,1,0)
+        self.relu14 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+        output = []
+        out = self.conv1(x)
+        out = self.reflecPad1(out)
+        out = self.conv2(out)
+        out = self.relu2(out)
+        output.append(out)
+
+        out = self.reflecPad3(out)
+        out = self.conv3(out)
+        out = self.relu3(out)
+        out = self.maxPool(out)
+        out = self.reflecPad4(out)
+        out = self.conv4(out)
+        out = self.relu4(out)
+        output.append(out)
+
+        out = self.reflecPad5(out)
+        out = self.conv5(out)
+        out = self.relu5(out)
+        out = self.maxPool2(out)
+        out = self.reflecPad6(out)
+        out = self.conv6(out)
+        out = self.relu6(out)
+        output.append(out)
+
+        out = self.reflecPad7(out)
+        out = self.conv7(out)
+        out = self.relu7(out)
+        out = self.reflecPad8(out)
+        out = self.conv8(out)
+        out = self.relu8(out)
+        out = self.reflecPad9(out)
+        out = self.conv9(out)
+        out = self.relu9(out)
+        out = self.maxPool3(out)
+        out = self.reflecPad10(out)
+        out = self.conv10(out)
+        out = self.relu10(out)
+        output.append(out)
+
+        out = self.reflecPad11(out)
+        out = self.conv11(out)
+        out = self.relu11(out)
+        out = self.reflecPad12(out)
+        out = self.conv12(out)
+        out = self.relu12(out)
+        out = self.reflecPad13(out)
+        out = self.conv13(out)
+        out = self.relu13(out)
+        out = self.maxPool4(out)
+        out = self.reflecPad14(out)
+        out = self.conv14(out)
+        out = self.relu14(out)
+
+        output.append(out)
+        return output
+
+class VGGLoss(nn.Module):
+    def __init__(self, model_path):
+        super(VGGLoss, self).__init__()
+        self.vgg = encoder5().cuda()
+        self.vgg.load_state_dict(torch.load(os.path.join(model_path, 'vgg_r51.pth')))
+        self.criterion = nn.MSELoss()
+        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
+
+    def forward(self, x, y):
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        for i in range(4):
+            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
+        return loss
+
+class GANLoss(nn.Module):
+    def __init__(self, gan_mode = 'hinge', target_real_label=1.0, target_fake_label=0.0,
+                 tensor=torch.cuda.FloatTensor):
+        super(GANLoss, self).__init__()
+        self.real_label = target_real_label
+        self.fake_label = target_fake_label
+        self.real_label_tensor = None
+        self.fake_label_tensor = None
+        self.zero_tensor = None
+        self.Tensor = tensor
+        self.gan_mode = gan_mode
+        if gan_mode == 'ls':
+            pass
+        elif gan_mode == 'original':
+            pass
+        elif gan_mode == 'w':
+            pass
+        elif gan_mode == 'hinge':
+            pass
+        else:
+            raise ValueError('Unexpected gan_mode {}'.format(gan_mode))
+
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            if self.real_label_tensor is None:
+                self.real_label_tensor = self.Tensor(1).fill_(self.real_label)
+                self.real_label_tensor.requires_grad_(False)
+            return self.real_label_tensor.expand_as(input)
+        else:
+            if self.fake_label_tensor is None:
+                self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)
+                self.fake_label_tensor.requires_grad_(False)
+            return self.fake_label_tensor.expand_as(input)
+
+    def get_zero_tensor(self, input):
+        if self.zero_tensor is None:
+            self.zero_tensor = self.Tensor(1).fill_(0)
+            self.zero_tensor.requires_grad_(False)
+        return self.zero_tensor.expand_as(input)
+
+    def loss(self, input, target_is_real, for_discriminator=True):
+        if self.gan_mode == 'original':  # cross entropy loss
+            target_tensor = self.get_target_tensor(input, target_is_real)
+            loss = F.binary_cross_entropy_with_logits(input, target_tensor)
+            return loss
+        elif self.gan_mode == 'ls':
+            target_tensor = self.get_target_tensor(input, target_is_real)
+            return F.mse_loss(input, target_tensor)
+        elif self.gan_mode == 'hinge':
+            if for_discriminator:
+                if target_is_real:
+                    minval = torch.min(input - 1, self.get_zero_tensor(input))
+                    loss = -torch.mean(minval)
+                else:
+                    minval = torch.min(-input - 1, self.get_zero_tensor(input))
+                    loss = -torch.mean(minval)
+            else:
+                assert target_is_real, "The generator's hinge loss must be aiming for real"
+                loss = -torch.mean(input)
+            return loss
+        else:
+            # wgan
+            if target_is_real:
+                return -input.mean()
+            else:
+                return input.mean()
+
+    def __call__(self, input, target_is_real, for_discriminator=True):
+        # computing loss is a bit complicated because |input| may not be
+        # a tensor, but list of tensors in case of multiscale discriminator
+        if isinstance(input, list):
+            loss = 0
+            for pred_i in input:
+                if isinstance(pred_i, list):
+                    pred_i = pred_i[-1]
+                loss_tensor = self.loss(pred_i, target_is_real, for_discriminator)
+                bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)
+                new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)
+                loss += new_loss
+            return loss / len(input)
+        else:
+            return self.loss(input, target_is_real, for_discriminator)
+
+class SPADE_LOSS(nn.Module):
+    def __init__(self, model_path, lambda_feat = 1):
+        super(SPADE_LOSS, self).__init__()
+        self.criterionVGG = VGGLoss(model_path)
+        self.criterionGAN = GANLoss('hinge')
+        self.criterionL1 = nn.L1Loss()
+        self.discriminator = NLayerDiscriminator()
+        self.lambda_feat = lambda_feat
+
+    def forward(self, x, y, for_discriminator = False):
+        pred_real = self.discriminator(y)
+        if not for_discriminator:
+            pred_fake = self.discriminator(x)
+            VGGLoss = self.criterionVGG(x, y)
+            GANLoss = self.criterionGAN(pred_fake, True, for_discriminator = False)
+
+            # feature matching loss
+            # last output is the final prediction, so we exclude it
+            num_intermediate_outputs = len(pred_fake) - 1
+            GAN_Feat_loss = 0
+            for j in range(num_intermediate_outputs):  # for each layer output
+                unweighted_loss = self.criterionL1(pred_fake[j], pred_real[j].detach())
+                GAN_Feat_loss += unweighted_loss * self.lambda_feat
+            L1Loss = self.criterionL1(x, y)
+            return VGGLoss, GANLoss, GAN_Feat_loss, L1Loss
+        else:
+            pred_fake = self.discriminator(x.detach())
+            GANLoss = self.criterionGAN(pred_fake, False, for_discriminator = True)
+            GANLoss += self.criterionGAN(pred_real, True, for_discriminator = True)
+            return GANLoss
+
+class ContrastiveLoss(nn.Module):
+    """
+    Contrastive loss
+    Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise
+    """
+
+    def __init__(self, margin):
+        super(ContrastiveLoss, self).__init__()
+        self.margin = margin
+        self.eps = 1e-9
+
+    def forward(self, out1, out2, target, size_average=True, norm = True):
+        if norm:
+            output1 = out1 / out1.pow(2).sum(1, keepdim=True).sqrt()
+            output2 = out1 / out2.pow(2).sum(1, keepdim=True).sqrt()
+        distances = (output2 - output1).pow(2).sum(1)  # squared distances
+        losses = 0.5 * (target.float() * distances +
+                        (1 + -1 * target).float() * F.relu(self.margin - (distances + self.eps).sqrt()).pow(2))
+        return losses.mean() if size_average else losses.sum()