First commit

README.md

@@ -1,8 +1,8 @@
 ---
 title: Instagram Filter Removal
-emoji: 🐨
-colorFrom: purple
-colorTo: pink
+emoji: 👀
+colorFrom: gray
+colorTo: green
 sdk: gradio
 app_file: app.py
 pinned: false
@@ -10,28 +10,28 @@ pinned: false
 
 # Configuration
 
-`title`: _string_
+`title`: _string_  
 Display title for the Space
 
-`emoji`: _string_
+`emoji`: _string_  
 Space emoji (emoji-only character allowed)
 
-`colorFrom`: _string_
+`colorFrom`: _string_  
 Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
 
-`colorTo`: _string_
+`colorTo`: _string_  
 Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
 
-`sdk`: _string_
-Can be either `gradio`, `streamlit`, or `static`
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
 
-`sdk_version` : _string_
+`sdk_version` : _string_  
 Only applicable for `streamlit` SDK.  
 See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
 
-`app_file`: _string_
-Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
 Path is relative to the root of the repository.
 
-`pinned`: _boolean_
+`pinned`: _boolean_  
 Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,79 @@
+import requests
+import os
+import gradio as gr
+import numpy as np
+import torch
+import torchvision.models as models
+
+from configs.default import get_cfg_defaults
+from modeling.build import build_model
+from utils.data_utils import linear_scaling
+
+
+url = "https://www.dropbox.com/s/y97z812sxa1kvrg/ifrnet.pth?dl=1"
+r = requests.get(url, stream=True)
+if not os.path.exists("ifrnet.pth"):
+    with open("ifrnet.pth", 'wb') as f:
+        for data in r:
+            f.write(data)
+
+cfg = get_cfg_defaults()
+cfg.MODEL.CKPT = "ifrnet.pth"
+net, _ = build_model(cfg)
+net = net.eval()
+vgg16 = models.vgg16(pretrained=True).features.eval()
+
+
+def load_checkpoints_from_ckpt(ckpt_path):
+    checkpoints = torch.load(ckpt_path, map_location=torch.device('cpu'))
+    net.load_state_dict(checkpoints["ifr"])
+
+
+load_checkpoints_from_ckpt(cfg.MODEL.CKPT)
+
+
+def filter_removal(img):
+    arr = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
+    arr = torch.tensor(arr).float() / 255.
+    arr = linear_scaling(arr)
+    with torch.no_grad():
+        feat = vgg16(arr)
+        out, _ = net(arr, feat)
+        out = torch.clamp(out, max=1., min=0.)
+        return out.squeeze(0).permute(1, 2, 0).numpy()
+
+
+title = "Instagram Filter Removal on Fashionable Images"
+description = "This is the demo for IFRNet, filter removal on fashionable images on Instagram. " \
+              "To use it, simply upload your filtered image, or click one of the examples to load them."
+article = "<p style='text-align: center'><a href='https://openaccess.thecvf.com/content/CVPR2021W/NTIRE/papers/Kinli_Instagram_Filter_Removal_on_Fashionable_Images_CVPRW_2021_paper.pdf'>Paper</a> | <a href='https://github.com/birdortyedi/instagram-filter-removal-pytorch'>Github</a></p>"
+
+gr.Interface(
+    filter_removal,
+    gr.inputs.Image(shape=(256, 256)),
+    gr.outputs.Image(),
+    title=title,
+    description=description,
+    article=article,
+    allow_flagging=False,
+    examples_per_page=17,
+    examples=[
+        ["images/examples/98_He-Fe.jpg"],
+        ["images/examples/2_Brannan.jpg"],
+        ["images/examples/12_Toaster.jpg"],
+        ["images/examples/18_Gingham.jpg"],
+        ["images/examples/11_Sutro.jpg"],
+        ["images/examples/9_Lo-Fi.jpg"],
+        ["images/examples/3_Mayfair.jpg"],
+        ["images/examples/4_Hudson.jpg"],
+        ["images/examples/5_Amaro.jpg"],
+        ["images/examples/6_1977.jpg"],
+        ["images/examples/8_Valencia.jpg"],
+        ["images/examples/16_Lo-Fi.jpg"],
+        ["images/examples/10_Nashville.jpg"],
+        ["images/examples/15_X-ProII.jpg"],
+        ["images/examples/14_Willow.jpg"],
+        ["images/examples/30_Perpetua.jpg"],
+        ["images/examples/1_Clarendon.jpg"],
+    ]
+).launch()

config.yaml

@@ -0,0 +1,98 @@
+wandb_version: 1
+
+DATASET:
+  desc: null
+  value:
+    MEAN:
+    - 0.5
+    - 0.5
+    - 0.5
+    NAME: IFFI
+    ROOT: ../../Downloads/IFFI-dataset/train
+    SIZE: 256
+    STD:
+    - 0.5
+    - 0.5
+    - 0.5
+    TEST_ROOT: ../../Downloads/IFFI-dataset/test
+MODEL:
+  desc: null
+  value:
+    D:
+      NAME: 1-ChOutputDiscriminator
+      NUM_CHANNELS: 32
+      NUM_CRITICS: 5
+      SOLVER:
+        BETAS:
+        - 0.5
+        - 0.9
+        DECAY_RATE: 0.5
+        LR: 0.001
+        SCHEDULER: []
+    IFR:
+      DESTYLER_CHANNELS: 32
+      NAME: InstaFilterRemovalNetwork
+      NUM_CHANNELS: 32
+      SOLVER:
+        BETAS:
+        - 0.5
+        - 0.9
+        DECAY_RATE: 0
+        LR: 0.0002
+        SCHEDULER: []
+    IS_TRAIN: true
+    NAME: ifrnet
+    NUM_CLASS: 17
+OPTIM:
+  desc: null
+  value:
+    ADVERSARIAL: 0.001
+    AUX: 0.5
+    GP: 10
+    MASK: 1
+    RECON: 1.4
+    SEMANTIC: 0.0001
+    TEXTURE: 0.001
+SYSTEM:
+  desc: null
+  value:
+    NUM_GPU: 2
+    NUM_WORKERS: 4
+TEST:
+  desc: null
+  value:
+    ABLATION: false
+    BATCH_SIZE: 64
+    IMG_ID: 52
+    OUTPUT_DIR: ./outputs
+    WEIGHTS: ''
+TRAIN:
+  desc: null
+  value:
+    BATCH_SIZE: 8
+    IS_TRAIN: true
+    LOG_INTERVAL: 100
+    NUM_TOTAL_STEP: 120000
+    RESUME: true
+    SAVE_DIR: ./weights
+    SAVE_INTERVAL: 1000
+    SHUFFLE: true
+    START_STEP: 0
+    TUNE: false
+    VISUALIZE_INTERVAL: 100
+WANDB:
+  desc: null
+  value:
+    ENTITY: vvgl-ozu
+    LOG_DIR: ./logs/ifrnet_IFFI_120000step_8bs_0.0002lr_2gpu_9run
+    NUM_ROW: 0
+    PROJECT_NAME: instagram-filter-removal
+    RUN: 9
+_wandb:
+  desc: null
+  value:
+    cli_version: 0.9.1
+    framework: torch
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    python_version: 3.6.9

configs/default.py

@@ -0,0 +1,88 @@
+from yacs.config import CfgNode as CN
+
+_C = CN()
+
+_C.SYSTEM = CN()
+_C.SYSTEM.NUM_GPU = 2
+_C.SYSTEM.NUM_WORKERS = 4
+
+_C.WANDB = CN()
+_C.WANDB.PROJECT_NAME = "instagram-filter-removal"
+_C.WANDB.ENTITY = "vvgl-ozu"
+_C.WANDB.RUN = 12
+_C.WANDB.LOG_DIR = ""
+_C.WANDB.NUM_ROW = 0
+
+_C.TRAIN = CN()
+_C.TRAIN.NUM_TOTAL_STEP = 120000
+_C.TRAIN.START_STEP = 0
+_C.TRAIN.BATCH_SIZE = 8
+_C.TRAIN.SHUFFLE = True
+_C.TRAIN.LOG_INTERVAL = 100
+_C.TRAIN.SAVE_INTERVAL = 5000
+_C.TRAIN.SAVE_DIR = "./weights"
+_C.TRAIN.RESUME = True
+_C.TRAIN.VISUALIZE_INTERVAL = 100
+_C.TRAIN.TUNE = False
+
+_C.MODEL = CN()
+_C.MODEL.NAME = "ifr-no-aux"
+_C.MODEL.IS_TRAIN = True
+_C.MODEL.NUM_CLASS = 17
+_C.MODEL.CKPT = ""
+
+_C.MODEL.IFR = CN()
+_C.MODEL.IFR.NAME = "InstaFilterRemovalNetwork"
+_C.MODEL.IFR.NUM_CHANNELS = 32
+_C.MODEL.IFR.DESTYLER_CHANNELS = 32
+_C.MODEL.IFR.SOLVER = CN()
+_C.MODEL.IFR.SOLVER.LR = 2e-4
+_C.MODEL.IFR.SOLVER.BETAS = (0.5, 0.9)
+_C.MODEL.IFR.SOLVER.SCHEDULER = []
+_C.MODEL.IFR.SOLVER.DECAY_RATE = 0.
+
+_C.MODEL.D = CN()
+_C.MODEL.D.NAME = "1-ChOutputDiscriminator"
+_C.MODEL.D.NUM_CHANNELS = 32
+_C.MODEL.D.NUM_CRITICS = 5
+_C.MODEL.D.SOLVER = CN()
+_C.MODEL.D.SOLVER.LR = 1e-3
+_C.MODEL.D.SOLVER.BETAS = (0.5, 0.9)
+_C.MODEL.D.SOLVER.SCHEDULER = []
+_C.MODEL.D.SOLVER.DECAY_RATE = 0.5
+
+_C.OPTIM = CN()
+_C.OPTIM.GP = 10
+_C.OPTIM.MASK = 1
+_C.OPTIM.RECON = 1.4
+_C.OPTIM.SEMANTIC = 1e-4
+_C.OPTIM.TEXTURE = 1e-3
+_C.OPTIM.ADVERSARIAL = 1e-3
+_C.OPTIM.AUX = 0.5
+
+_C.DATASET = CN()
+_C.DATASET.NAME = "IFFI"  # "IFFI"  # "DIV2K?" #
+_C.DATASET.ROOT = "../../Datasets/IFFI-dataset/train"  # "../../Datasets/IFFI-dataset"  # "/media/birdortyedi/e5042b8f-ca5e-4a22-ac68-7e69ff648bc4/IFFI-dataset"
+_C.DATASET.TEST_ROOT = "../../Datasets/IFFI-dataset"
+_C.DATASET.SIZE = 256
+_C.DATASET.CROP_SIZE = 512
+_C.DATASET.MEAN = [0.5, 0.5, 0.5]
+_C.DATASET.STD = [0.5, 0.5, 0.5]
+
+_C.TEST = CN()
+_C.TEST.OUTPUT_DIR = "./outputs"
+_C.TEST.ABLATION = False
+_C.TEST.WEIGHTS = ""
+_C.TEST.BATCH_SIZE = 64
+_C.TEST.IMG_ID = 52
+
+
+def get_cfg_defaults():
+    """Get a yacs CfgNode object with default values for my_project."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    return _C.clone()
+
+
+# provide a way to import the defaults as a global singleton:
+cfg = _C  # users can `from config import cfg`

modeling/base.py

@@ -0,0 +1,60 @@
+from torch import nn
+
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def forward(self, x, y):
+        pass
+
+    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 set_requires_grad(self, requires_grad=False):
+        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
+        Parameters:
+        requires_grad (bool) -- whether the networks require gradients or not
+        """
+        for param in self.parameters():
+            param.requires_grad = requires_grad
+
+    def init_weights(self, init_type='xavier', 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:
+                    nn.init.normal_(m.weight.data, 1.0, gain)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    nn.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':
+                    nn.init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == 'xavier':
+                    nn.init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == 'xavier_uniform':
+                    nn.init.xavier_uniform_(m.weight.data, gain=1.0)
+                elif init_type == 'kaiming':
+                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                elif init_type == 'orthogonal':
+                    nn.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:
+                    nn.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)

modeling/benchmark.py

@@ -0,0 +1,62 @@
+import torch
+from torch import nn
+
+from modeling.base import BaseNetwork
+from modeling.ifrnet import Flatten
+from modules.blocks import DestyleResBlock, Destyler, ResBlock
+
+
+class UNet(BaseNetwork):
+    def __init__(self, base_n_channels):
+        super(UNet, self).__init__()
+
+        self.ds_res1 = ResBlock(channels_in=3, channels_out=base_n_channels, kernel_size=5, stride=1, padding=2)
+        self.ds_res2 = ResBlock(channels_in=base_n_channels, channels_out=base_n_channels * 2, kernel_size=3, stride=2, padding=1)
+        self.ds_res3 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.ds_res4 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 4, kernel_size=3, stride=2, padding=1)
+        self.ds_res5 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.ds_res6 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 8, kernel_size=3, stride=2, padding=1)
+
+        self.upsample = nn.UpsamplingNearest2d(scale_factor=2.0)
+
+        self.res1 = ResBlock(channels_in=base_n_channels * 8, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.res2 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.res3 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.res4 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.res5 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels, kernel_size=3, stride=1, padding=1)
+
+        self.conv1 = nn.Conv2d(base_n_channels, 3, kernel_size=3, stride=1, padding=1)
+
+        self.init_weights(init_type="normal", gain=0.02)
+
+    def forward(self, x):
+        out = self.ds_res1(x)
+        out = self.ds_res2(out)
+        out = self.ds_res3(out)
+        out = self.ds_res4(out)
+        out = self.ds_res5(out)
+        aux = self.ds_res6(out)
+
+        out = self.upsample(aux)
+        out = self.res1(out)
+        out = self.res2(out)
+        out = self.upsample(out)
+        out = self.res3(out)
+        out = self.res4(out)
+        out = self.upsample(out)
+        out = self.res5(out)
+        out = self.conv1(out)
+
+        return out, aux
+
+
+if __name__ == '__main__':
+    import torchvision
+    x = torch.rand((2, 3, 256, 256)).cuda()
+    unet = UNet(32, 32).cuda()
+    vgg16 = torchvision.models.vgg16(pretrained=True).features.eval().cuda()
+    with torch.no_grad():
+        vgg_feat = vgg16(x)
+    out = unet(x, vgg_feat)
+
+    print(out.size())

modeling/build.py

@@ -0,0 +1,19 @@
+from modeling.ifrnet import IFRNet, Discriminator, PatchDiscriminator, MLP
+from modeling.benchmark import UNet
+
+
+def build_model(args):
+    if args.MODEL.NAME.lower() == "ifrnet":
+        net = IFRNet(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, destyler_n_channels=args.MODEL.IFR.DESTYLER_CHANNELS)
+        mlp = MLP(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, num_class=args.MODEL.NUM_CLASS)
+    elif args.MODEL.NAME.lower() == "ifr-no-aux":
+        net = IFRNet(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, destyler_n_channels=args.MODEL.IFR.DESTYLER_CHANNELS)
+        mlp = None
+    else:
+        raise NotImplementedError
+    return net, mlp
+
+
+def build_discriminators(args):
+    return Discriminator(base_n_channels=args.MODEL.D.NUM_CHANNELS), PatchDiscriminator(base_n_channels=args.MODEL.D.NUM_CHANNELS)
+

modeling/ifrnet.py

@@ -0,0 +1,166 @@
+import torch
+from torch import nn
+from torch.nn.utils import spectral_norm
+
+from modeling.base import BaseNetwork
+from modules.blocks import DestyleResBlock, Destyler, ResBlock
+
+
+class IFRNet(BaseNetwork):
+    def __init__(self, base_n_channels, destyler_n_channels):
+        super(IFRNet, self).__init__()
+        self.destyler = Destyler(in_features=32768, num_features=destyler_n_channels)  # from vgg features
+
+        self.ds_fc1 = nn.Linear(destyler_n_channels, base_n_channels * 2)
+        self.ds_res1 = DestyleResBlock(channels_in=3, channels_out=base_n_channels, kernel_size=5, stride=1, padding=2)
+        self.ds_fc2 = nn.Linear(destyler_n_channels, base_n_channels * 4)
+        self.ds_res2 = DestyleResBlock(channels_in=base_n_channels, channels_out=base_n_channels * 2, kernel_size=3, stride=2, padding=1)
+        self.ds_fc3 = nn.Linear(destyler_n_channels, base_n_channels * 4)
+        self.ds_res3 = DestyleResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.ds_fc4 = nn.Linear(destyler_n_channels, base_n_channels * 8)
+        self.ds_res4 = DestyleResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 4, kernel_size=3, stride=2, padding=1)
+        self.ds_fc5 = nn.Linear(destyler_n_channels, base_n_channels * 8)
+        self.ds_res5 = DestyleResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.ds_fc6 = nn.Linear(destyler_n_channels, base_n_channels * 16)
+        self.ds_res6 = DestyleResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 8, kernel_size=3, stride=2, padding=1)
+
+        self.upsample = nn.UpsamplingNearest2d(scale_factor=2.0)
+
+        self.res1 = ResBlock(channels_in=base_n_channels * 8, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.res2 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 4, kernel_size=3, stride=1, padding=1)
+        self.res3 = ResBlock(channels_in=base_n_channels * 4, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.res4 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels * 2, kernel_size=3, stride=1, padding=1)
+        self.res5 = ResBlock(channels_in=base_n_channels * 2, channels_out=base_n_channels, kernel_size=3, stride=1, padding=1)
+
+        self.conv1 = nn.Conv2d(base_n_channels, 3, kernel_size=3, stride=1, padding=1)
+
+        self.init_weights(init_type="normal", gain=0.02)
+
+    def forward(self, x, vgg_feat):
+        b_size, ch, h, w = vgg_feat.size()
+        vgg_feat = vgg_feat.view(b_size, ch * h * w)
+        vgg_feat = self.destyler(vgg_feat)
+
+        out = self.ds_res1(x, self.ds_fc1(vgg_feat))
+        out = self.ds_res2(out, self.ds_fc2(vgg_feat))
+        out = self.ds_res3(out, self.ds_fc3(vgg_feat))
+        out = self.ds_res4(out, self.ds_fc4(vgg_feat))
+        out = self.ds_res5(out, self.ds_fc5(vgg_feat))
+        aux = self.ds_res6(out, self.ds_fc6(vgg_feat))
+
+        out = self.upsample(aux)
+        out = self.res1(out)
+        out = self.res2(out)
+        out = self.upsample(out)
+        out = self.res3(out)
+        out = self.res4(out)
+        out = self.upsample(out)
+        out = self.res5(out)
+        out = self.conv1(out)
+
+        return out, aux
+
+
+class MLP(nn.Module):
+    def __init__(self, base_n_channels, num_class=14):
+        super(MLP, self).__init__()
+        self.aux_classifier = nn.Sequential(
+            nn.Conv2d(base_n_channels * 8, base_n_channels * 4, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(2),
+            nn.Conv2d(base_n_channels * 4, base_n_channels * 2, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(2),
+            # nn.Conv2d(base_n_channels * 2, base_n_channels * 1, kernel_size=3, stride=1, padding=1),
+            # nn.MaxPool2d(2),
+            Flatten(),
+            nn.Linear(base_n_channels * 8 * 8 * 2, num_class),
+            # nn.Softmax(dim=-1)
+        )
+
+    def forward(self, x):
+        return self.aux_classifier(x)
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        """
+        Note that input.size(0) is usually the batch size.
+        So what it does is that given any input with input.size(0) # of batches,
+        will flatten to be 1 * nb_elements.
+        """
+        batch_size = input.size(0)
+        out = input.view(batch_size, -1)
+        return out  # (batch_size, *size)
+
+
+class Discriminator(BaseNetwork):
+    def __init__(self, base_n_channels):
+        """
+        img_size : (int, int, int)
+            Height and width must be powers of 2.  E.g. (32, 32, 1) or
+            (64, 128, 3). Last number indicates number of channels, e.g. 1 for
+            grayscale or 3 for RGB
+        """
+        super(Discriminator, self).__init__()
+
+        self.image_to_features = nn.Sequential(
+            spectral_norm(nn.Conv2d(3, base_n_channels, 5, 2, 2)),
+            nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(nn.Conv2d(base_n_channels, 2 * base_n_channels, 5, 2, 2)),
+            nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(nn.Conv2d(2 * base_n_channels, 2 * base_n_channels, 5, 2, 2)),
+            nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(nn.Conv2d(2 * base_n_channels, 4 * base_n_channels, 5, 2, 2)),
+            nn.LeakyReLU(0.2, inplace=True),
+            # spectral_norm(nn.Conv2d(4 * base_n_channels, 4 * base_n_channels, 5, 2, 2)),
+            # nn.LeakyReLU(0.2, inplace=True),
+            spectral_norm(nn.Conv2d(4 * base_n_channels, 8 * base_n_channels, 5, 1, 1)),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+
+        output_size = 8 * base_n_channels * 3 * 3
+        self.features_to_prob = nn.Sequential(
+            spectral_norm(nn.Conv2d(8 * base_n_channels, 2 * base_n_channels, 5, 2, 1)),
+            Flatten(),
+            nn.Linear(output_size, 1)
+        )
+
+        self.init_weights(init_type="normal", gain=0.02)
+
+    def forward(self, input_data):
+        x = self.image_to_features(input_data)
+        return self.features_to_prob(x)
+
+
+class PatchDiscriminator(Discriminator):
+    def __init__(self, base_n_channels):
+        super(PatchDiscriminator, self).__init__(base_n_channels)
+
+        self.features_to_prob = nn.Sequential(
+            spectral_norm(nn.Conv2d(8 * base_n_channels, 1, 1)),
+            Flatten()
+        )
+
+    def forward(self, input_data):
+        x = self.image_to_features(input_data)
+        return self.features_to_prob(x)
+
+
+if __name__ == '__main__':
+    import torchvision
+    ifrnet = IFRNet(32, 128).cuda()
+    x = torch.rand((2, 3, 256, 256)).cuda()
+    vgg16 = torchvision.models.vgg16(pretrained=True).features.eval().cuda()
+    with torch.no_grad():
+        vgg_feat = vgg16(x)
+    output, aux_out = ifrnet(x, vgg_feat)
+    print(output.size())
+    print(aux_out.size())
+
+    disc = Discriminator(32).cuda()
+    d_out = disc(output)
+    print(d_out.size())
+
+    patch_disc = PatchDiscriminator(32).cuda()
+    p_d_out = patch_disc(output)
+    print(p_d_out.size())
+

modules/blocks.py

@@ -0,0 +1,93 @@
+from torch import nn
+
+from modules.normalization import AdaIN
+
+
+class DestyleResBlock(nn.Module):
+    def __init__(self, channels_out, kernel_size, channels_in=None, stride=1, dilation=1, padding=1, use_dropout=False):
+        super(DestyleResBlock, self).__init__()
+
+        # uses 1x1 convolutions for downsampling
+        if not channels_in or channels_in == channels_out:
+            channels_in = channels_out
+            self.projection = None
+        else:
+            self.projection = nn.Conv2d(channels_in, channels_out, kernel_size=1, stride=stride, dilation=1)
+        self.use_dropout = use_dropout
+
+        self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation)
+        self.lrelu1 = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=kernel_size, stride=1, padding=padding, dilation=dilation)
+        self.adain = AdaIN()
+        if self.use_dropout:
+            self.dropout = nn.Dropout()
+        self.lrelu2 = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x, feat):
+        residual = x
+        out = self.conv1(x)
+        out = self.lrelu1(out)
+        out = self.conv2(out)
+        _, _, h, w = out.size()
+        out = self.adain(out, feat)
+        if self.use_dropout:
+            out = self.dropout(out)
+        if self.projection:
+            residual = self.projection(x)
+        out = out + residual
+        out = self.lrelu2(out)
+        return out
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channels_out, kernel_size, channels_in=None, stride=1, dilation=1, padding=1, use_dropout=False):
+        super(ResBlock, self).__init__()
+
+        # uses 1x1 convolutions for downsampling
+        if not channels_in or channels_in == channels_out:
+            channels_in = channels_out
+            self.projection = None
+        else:
+            self.projection = nn.Conv2d(channels_in, channels_out, kernel_size=1, stride=stride, dilation=1)
+        self.use_dropout = use_dropout
+
+        self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation)
+        self.lrelu1 = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=kernel_size, stride=1, padding=padding, dilation=dilation)
+        self.n2 = nn.BatchNorm2d(channels_out)
+        if self.use_dropout:
+            self.dropout = nn.Dropout()
+        self.lrelu2 = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.lrelu1(out)
+        out = self.conv2(out)
+        # out = self.n2(out)
+        if self.use_dropout:
+            out = self.dropout(out)
+        if self.projection:
+            residual = self.projection(x)
+        out = out + residual
+        out = self.lrelu2(out)
+        return out
+
+
+class Destyler(nn.Module):
+    def __init__(self, in_features, num_features):
+        super(Destyler, self).__init__()
+        self.fc1 = nn.Linear(in_features, num_features)
+        self.fc2 = nn.Linear(num_features, num_features)
+        self.fc3 = nn.Linear(num_features, num_features)
+        self.fc4 = nn.Linear(num_features, num_features)
+        self.fc5 = nn.Linear(num_features, num_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        x = self.fc4(x)
+        x = self.fc5(x)
+        return x
+

modules/normalization.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+
+class AdaIN(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        ch = y.size(1)
+        sigma, mu = torch.split(y.unsqueeze(-1).unsqueeze(-1), [ch // 2, ch // 2], dim=1)
+
+        x_mu = x.mean(dim=[2, 3], keepdim=True)
+        x_sigma = x.std(dim=[2, 3], keepdim=True)
+
+        return sigma * ((x - x_mu) / x_sigma) + mu

requirements.txt

@@ -0,0 +1,12 @@
+numpy>=1.17.0
+requests>=2.25.1
+torchvision>=0.6.0
+yacs>=0.1.7
+kornia>=0.3.1
+matplotlib>=3.3.4
+torch>=1.5.0
+glog>=0.3.1
+gradio>=1.6.4
+seaborn>=0.11.0
+Pillow>=8.2.0
+scikit_learn>=0.24.1

utils/data_utils.py

@@ -0,0 +1,6 @@
+def linear_scaling(x):
+    return (x * 255.) / 127.5 - 1.
+
+
+def linear_unscaling(x):
+    return (x + 1.) * 127.5 / 255.