Add application file

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/uxvax5sjx5iysyl/cifr.pth?dl=0"
+r = requests.get(url, stream=True)
+if not os.path.exists("cifr.pth"):
+    with open("cifr.pth", 'wb') as f:
+        for data in r:
+            f.write(data)
+
+cfg = get_cfg_defaults()
+cfg.MODEL.CKPT = "cifr.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('cuda'))
+    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 = "Contrastive Instagram Filter Removal (CIFR)"
+description = "This is the demo for CIFR, contrastive strategy for 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://arxiv.org/abs/2204.07486'>Contrastive Instagram Filter Removal (CIFR)</a> | <a href='https://github.com/birdortyedi/cifr-pytorch'>Github Repo</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()

configs/default.py

@@ -0,0 +1,114 @@
+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 = "contrastive-style-learning-for-ifr"
+_C.WANDB.ENTITY = "vvgl-ozu"
+_C.WANDB.RUN = 3
+_C.WANDB.LOG_DIR = ""
+_C.WANDB.NUM_ROW = 0
+
+_C.TRAIN = CN()
+_C.TRAIN.NUM_TOTAL_STEP = 200000
+_C.TRAIN.START_STEP = 0
+_C.TRAIN.BATCH_SIZE = 16
+_C.TRAIN.SHUFFLE = True
+_C.TRAIN.LOG_INTERVAL = 100
+_C.TRAIN.EVAL_INTERVAL = 1000
+_C.TRAIN.SAVE_INTERVAL = 1000
+_C.TRAIN.SAVE_DIR = "./weights"
+_C.TRAIN.RESUME = True
+_C.TRAIN.VISUALIZE_INTERVAL = 100
+_C.TRAIN.TUNE = False
+
+_C.MODEL = CN()
+_C.MODEL.NAME = "cifr"
+_C.MODEL.IS_TRAIN = True
+_C.MODEL.NUM_CLASS = 17
+_C.MODEL.CKPT = ""
+_C.MODEL.PRETRAINED = ""
+
+_C.MODEL.IFR = CN()
+_C.MODEL.IFR.NAME = "ContrastiveInstaFilterRemovalNetwork"
+_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.999)
+_C.MODEL.IFR.SOLVER.SCHEDULER = []
+_C.MODEL.IFR.SOLVER.DECAY_RATE = 0.
+_C.MODEL.IFR.DS_FACTOR = 4
+
+_C.MODEL.PATCH = CN()
+_C.MODEL.PATCH.NUM_CHANNELS = 256
+_C.MODEL.PATCH.NUM_PATCHES = 256
+_C.MODEL.PATCH.NUM_LAYERS = 6
+_C.MODEL.PATCH.USE_MLP = True
+_C.MODEL.PATCH.SHUFFLE_Y = True
+_C.MODEL.PATCH.LR = 1e-4
+_C.MODEL.PATCH.BETAS = (0.5, 0.999)
+_C.MODEL.PATCH.T = 0.07
+
+_C.MODEL.D = CN()
+_C.MODEL.D.NAME = "1-ChOutputDiscriminator"
+_C.MODEL.D.NUM_CHANNELS = 32
+_C.MODEL.D.NUM_CRITICS = 3
+_C.MODEL.D.SOLVER = CN()
+_C.MODEL.D.SOLVER.LR = 1e-4
+_C.MODEL.D.SOLVER.BETAS = (0.5, 0.999)
+_C.MODEL.D.SOLVER.SCHEDULER = []
+_C.MODEL.D.SOLVER.DECAY_RATE = 0.01
+
+_C.ESRGAN = CN()
+_C.ESRGAN.WEIGHTS = "weights/RealESRGAN_x{}plus.pth"
+
+_C.FASHIONMASKRCNN = CN()
+_C.FASHIONMASKRCNN.CFG_PATH = "configs/fashion.yaml"
+_C.FASHIONMASKRCNN.WEIGHTS = "weights/fashion.pth"
+_C.FASHIONMASKRCNN.SCORE_THRESH_TEST = 0.6
+_C.FASHIONMASKRCNN.MIN_SIZE_TEST = 512
+
+_C.OPTIM = CN()
+_C.OPTIM.GP = 10.
+_C.OPTIM.MASK = 1
+_C.OPTIM.RECON = 1.4
+_C.OPTIM.SEMANTIC = 1e-1
+_C.OPTIM.TEXTURE = 2e-1
+_C.OPTIM.ADVERSARIAL = 1e-3
+_C.OPTIM.AUX = 0.5
+_C.OPTIM.CONTRASTIVE = 0.1
+_C.OPTIM.NLL = 1.0
+
+_C.DATASET = CN()
+_C.DATASET.NAME = "IFFI"
+_C.DATASET.ROOT = "../../Downloads/IFFI-dataset/train"  # "../../Downloads/IFFI-dataset/train"
+_C.DATASET.TEST_ROOT = "../../Datasets/IFFI-dataset/test"  # "../../Downloads/IFFI-dataset/test"
+_C.DATASET.DS_TEST_ROOT = "../../Downloads/IFFI-dataset/test/"  # "../../Downloads/IFFI-dataset/test"
+_C.DATASET.DS_JSON_FILE = "../../Downloads/IFFI-dataset-only-orgs/instances_default.json"
+_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 = 32
+_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`

layers/blocks.py

@@ -0,0 +1,93 @@
+from torch import nn
+
+from layers.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
+

layers/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

modeling/arch.py

@@ -0,0 +1,272 @@
+import torch
+from torch import nn
+from torch.nn.utils import spectral_norm
+
+from modeling.base import BaseNetwork
+from layers.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 CIFR_Encoder(IFRNet):
+    def __init__(self, base_n_channels, destyler_n_channels):
+        super(CIFR_Encoder, self).__init__(base_n_channels, destyler_n_channels)
+
+    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)
+
+        feat1 = self.ds_res1(x, self.ds_fc1(vgg_feat))
+        feat2 = self.ds_res2(feat1, self.ds_fc2(vgg_feat))
+        feat3 = self.ds_res3(feat2, self.ds_fc3(vgg_feat))
+        feat4 = self.ds_res4(feat3, self.ds_fc4(vgg_feat))
+        feat5 = self.ds_res5(feat4, self.ds_fc5(vgg_feat))
+        feat6 = self.ds_res6(feat5, self.ds_fc6(vgg_feat))
+
+        feats = [feat1, feat2, feat3, feat4, feat5, feat6]
+
+        out = self.upsample(feat6)
+        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, feats
+
+
+class Normalize(nn.Module):
+    def __init__(self, power=2):
+        super(Normalize, self).__init__()
+        self.power = power
+
+    def forward(self, x):
+        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)
+        out = x.div(norm + 1e-7)
+        return out
+
+
+class PatchSampleF(BaseNetwork):
+    def __init__(self, base_n_channels, style_or_content, use_mlp=False, nc=256):
+        # potential issues: currently, we use the same patch_ids for multiple images in the batch
+        super(PatchSampleF, self).__init__()
+        self.is_content = True if style_or_content == "content" else False
+        self.l2norm = Normalize(2)
+        self.use_mlp = use_mlp
+        self.nc = nc  # hard-coded
+
+        self.mlp_0 = nn.Sequential(*[nn.Linear(base_n_channels, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.mlp_1 = nn.Sequential(*[nn.Linear(base_n_channels * 2, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.mlp_2 = nn.Sequential(*[nn.Linear(base_n_channels * 2, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.mlp_3 = nn.Sequential(*[nn.Linear(base_n_channels * 4, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.mlp_4 = nn.Sequential(*[nn.Linear(base_n_channels * 4, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.mlp_5 = nn.Sequential(*[nn.Linear(base_n_channels * 8, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]).cuda()
+        self.init_weights(init_type="normal", gain=0.02)
+
+    @staticmethod
+    def gram_matrix(x):
+        # a, b, c, d = x.size()  # a=batch size(=1)
+        a, b = x.size()
+        # b=number of feature maps
+        # (c,d)=dimensions of a f. map (N=c*d)
+
+        # features = x.view(a * b, c * d)  # resise F_XL into \hat F_XL
+
+        G = torch.mm(x, x.t())  # compute the gram product
+
+        # we 'normalize' the values of the gram matrix
+        # by dividing by the number of element in each feature maps.
+        return G.div(a * b)
+
+    def forward(self, feats, num_patches=64, patch_ids=None):
+        return_ids = []
+        return_feats = []
+
+        for feat_id, feat in enumerate(feats):
+            B, C, H, W = feat.shape
+            feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2)
+            if num_patches > 0:
+                if patch_ids is not None:
+                    patch_id = patch_ids[feat_id]
+                else:
+                    patch_id = torch.randperm(feat_reshape.shape[1], device=feats[0].device)
+                    patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))]  # .to(patch_ids.device)
+                x_sample = feat_reshape[:, patch_id, :].flatten(0, 1)  # reshape(-1, x.shape[1])
+            else:
+                x_sample = feat_reshape
+                patch_id = []
+            if self.use_mlp:
+                mlp = getattr(self, 'mlp_%d' % feat_id)
+                x_sample = mlp(x_sample)
+            if not self.is_content:
+                x_sample = self.gram_matrix(x_sample)
+            return_ids.append(patch_id)
+            x_sample = self.l2norm(x_sample)
+
+            if num_patches == 0:
+                x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W])
+            return_feats.append(x_sample)
+        return return_feats, return_ids
+
+
+class MLP(nn.Module):
+    def __init__(self, base_n_channels, out_features=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, out_features),
+            # 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 = CIFR_Encoder(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, feats = ifrnet(x, vgg_feat)
+    print(output.size())
+    for i, feat in enumerate(feats):
+        print(i, feat.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())
+

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/build.py

@@ -0,0 +1,25 @@
+from modeling.arch import IFRNet, CIFR_Encoder, Discriminator, PatchDiscriminator, MLP, PatchSampleF
+
+
+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, out_features=args.MODEL.NUM_CLASS)
+    elif args.MODEL.NAME.lower() == "cifr":
+        net = CIFR_Encoder(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, destyler_n_channels=args.MODEL.IFR.DESTYLER_CHANNELS)
+        mlp = None
+    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)
+
+
+def build_patch_sampler(args):
+    return PatchSampleF(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, style_or_content="content", use_mlp=args.MODEL.PATCH.USE_MLP, nc=args.MODEL.PATCH.NUM_CHANNELS), \
+           PatchSampleF(base_n_channels=args.MODEL.IFR.NUM_CHANNELS, style_or_content="style", use_mlp=args.MODEL.PATCH.USE_MLP, nc=args.MODEL.PATCH.NUM_CHANNELS)

requirements.txt

@@ -0,0 +1,6 @@
+gradio==2.9.4
+numpy==1.21.2
+requests==2.27.1
+torch==1.10.1
+torchvision==0.11.2
+yacs==0.1.8

utils/data_utils.py

@@ -0,0 +1,29 @@
+import json
+
+
+def linear_scaling(x):
+    return (x * 255.) / 127.5 - 1.
+
+
+def linear_unscaling(x):
+    return (x + 1.) * 127.5 / 255.
+
+
+def read_json(path):
+    """
+    :param path (str or os.Path): JSON file path.
+    :return: (Dict): the data in the JSON file.
+    """
+    with open(path) as f:
+        data = json.load(f)
+    return data
+
+
+def write_json(path, datagroup):
+    """
+    :param path (str or os.Path): File path for the output JSON file.
+    :param datagroup (Dict): The data which should be dump to the JSON file.
+    :return: void.
+    """
+    with open(path, "w+") as f:
+        json.dump(datagroup, f)