initial commit

.gitignore

@@ -0,0 +1,14 @@
+env/
+.idea
+
+# Models
+pretrained_model/*.pth
+
+# Output
+test_img/*
+test_output/*
+
+.history/
+.vscode/
+__pycache__/
+.DS_Store

README.md

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

app.py

@@ -0,0 +1,74 @@
+import os
+
+os.system("pip install gradio==2.4.6")
+import torch
+import gradio as gr
+import numpy as np
+import torchvision.utils as vutils
+import torchvision.transforms as transforms
+
+from PIL import Image
+from torch.autograd import Variable
+from network.Transformer import Transformer
+
+LOAD_SIZE = 1280
+STYLE = "Shinkai"
+MODEL_PATH = "pretrained_model"
+COLOUR_MODEL = "RGB"
+
+model = Transformer()
+model.load_state_dict(
+    torch.load(os.path.join(MODEL_PATH, f"{STYLE}_net_G_float.pth"))
+)
+model.eval()
+
+disable_gpu = torch.cuda.is_available()
+
+
+def inference(img):
+    # load image
+    input_image = img.convert(COLOUR_MODEL)
+    input_image = np.asarray(input_image)
+    # RGB -> BGR
+    input_image = input_image[:, :, [2, 1, 0]]
+    input_image = transforms.ToTensor()(input_image).unsqueeze(0)
+    # preprocess, (-1, 1)
+    input_image = -1 + 2 * input_image
+
+    if disable_gpu:
+        input_image = Variable(input_image).float()
+    else:
+        input_image = Variable(input_image).cuda()
+
+    # forward
+    output_image = model(input_image)
+    output_image = output_image[0]
+    # BGR -> RGB
+    output_image = output_image[[2, 1, 0], :, :]
+    output_image = output_image.data.cpu().float() * 0.5 + 0.5
+
+    return output_image
+
+
+title = "AnimeBackgroundGAN"
+description = "CartoonGAN from [Chen et.al](http://openaccess.thecvf.com/content_cvpr_2018/CameraReady/2205.pdf) based on [Yijunmaverick's implementation](https://github.com/Yijunmaverick/CartoonGAN-Test-Pytorch-Torch)"
+article = "<p style='text-align: center'><a href='https://github.com/venture-anime/cartoongan-pytorch' target='_blank'>Github Repo</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=akiyamasho' alt='visitor badge'></center></p>"
+
+examples = [
+    ["examples/garden_in.jpeg", "examples/garden_out.jpg"],
+    ["examples/library_in.jpeg", "examples/library_out.jpg"],
+]
+
+
+gr.Interface(
+    inference,
+    [gr.inputs.Image(type="pil")],
+    gr.outputs.Image(type="pil"),
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    allow_flagging=False,
+    allow_screenshot=False,
+    enable_queue=True,
+).launch()

main.py

@@ -0,0 +1,76 @@
+import torch
+import os
+import numpy as np
+import torchvision.utils as vutils
+
+from PIL import Image
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+
+from network.Transformer import Transformer
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--input_dir", default="test_img")
+parser.add_argument("--load_size", default=1280)
+parser.add_argument("--model_path", default="./pretrained_model")
+parser.add_argument("--style", default="Shinkai")
+parser.add_argument("--output_dir", default="test_output")
+parser.add_argument("--gpu", type=int, default=0)
+
+opt = parser.parse_args()
+
+valid_ext = [".jpg", ".png"]
+
+# setup
+if not os.path.exists(opt.input_dir):
+    os.makedirs(opt.input_dir)
+if not os.path.exists(opt.output_dir):
+    os.makedirs(opt.output_dir)
+
+# load pretrained model
+model = Transformer()
+model.load_state_dict(
+    torch.load(os.path.join(opt.model_path, opt.style + "_net_G_float.pth"))
+)
+model.eval()
+
+disable_gpu = opt.gpu == -1 or not torch.cuda.is_available()
+
+if disable_gpu:
+    print("CPU mode")
+    model.float()
+else:
+    print("GPU mode")
+    model.cuda()
+
+for files in os.listdir(opt.input_dir):
+    ext = os.path.splitext(files)[1]
+    if ext not in valid_ext:
+        continue
+    # load image
+    input_image = Image.open(os.path.join(opt.input_dir, files)).convert("RGB")
+    input_image = np.asarray(input_image)
+    # RGB -> BGR
+    input_image = input_image[:, :, [2, 1, 0]]
+    input_image = transforms.ToTensor()(input_image).unsqueeze(0)
+    # preprocess, (-1, 1)
+    input_image = -1 + 2 * input_image
+    if disable_gpu:
+        input_image = Variable(input_image).float()
+    else:
+        input_image = Variable(input_image).cuda()
+
+    # forward
+    output_image = model(input_image)
+    output_image = output_image[0]
+    # BGR -> RGB
+    output_image = output_image[[2, 1, 0], :, :]
+    output_image = output_image.data.cpu().float() * 0.5 + 0.5
+    # save
+    vutils.save_image(
+        output_image,
+        os.path.join(opt.output_dir, files[:-4] + "_" + opt.style + ".jpg"),
+    )
+
+print("Done!")

network/Transformer.py

@@ -0,0 +1,180 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Transformer(nn.Module):
+    def __init__(self):
+        super(Transformer, self).__init__()
+        #
+        self.refpad01_1 = nn.ReflectionPad2d(3)
+        self.conv01_1 = nn.Conv2d(3, 64, 7)
+        self.in01_1 = InstanceNormalization(64)
+        # relu
+        self.conv02_1 = nn.Conv2d(64, 128, 3, 2, 1)
+        self.conv02_2 = nn.Conv2d(128, 128, 3, 1, 1)
+        self.in02_1 = InstanceNormalization(128)
+        # relu
+        self.conv03_1 = nn.Conv2d(128, 256, 3, 2, 1)
+        self.conv03_2 = nn.Conv2d(256, 256, 3, 1, 1)
+        self.in03_1 = InstanceNormalization(256)
+        # relu
+
+        ## res block 1
+        self.refpad04_1 = nn.ReflectionPad2d(1)
+        self.conv04_1 = nn.Conv2d(256, 256, 3)
+        self.in04_1 = InstanceNormalization(256)
+        # relu
+        self.refpad04_2 = nn.ReflectionPad2d(1)
+        self.conv04_2 = nn.Conv2d(256, 256, 3)
+        self.in04_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 2
+        self.refpad05_1 = nn.ReflectionPad2d(1)
+        self.conv05_1 = nn.Conv2d(256, 256, 3)
+        self.in05_1 = InstanceNormalization(256)
+        # relu
+        self.refpad05_2 = nn.ReflectionPad2d(1)
+        self.conv05_2 = nn.Conv2d(256, 256, 3)
+        self.in05_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 3
+        self.refpad06_1 = nn.ReflectionPad2d(1)
+        self.conv06_1 = nn.Conv2d(256, 256, 3)
+        self.in06_1 = InstanceNormalization(256)
+        # relu
+        self.refpad06_2 = nn.ReflectionPad2d(1)
+        self.conv06_2 = nn.Conv2d(256, 256, 3)
+        self.in06_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 4
+        self.refpad07_1 = nn.ReflectionPad2d(1)
+        self.conv07_1 = nn.Conv2d(256, 256, 3)
+        self.in07_1 = InstanceNormalization(256)
+        # relu
+        self.refpad07_2 = nn.ReflectionPad2d(1)
+        self.conv07_2 = nn.Conv2d(256, 256, 3)
+        self.in07_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 5
+        self.refpad08_1 = nn.ReflectionPad2d(1)
+        self.conv08_1 = nn.Conv2d(256, 256, 3)
+        self.in08_1 = InstanceNormalization(256)
+        # relu
+        self.refpad08_2 = nn.ReflectionPad2d(1)
+        self.conv08_2 = nn.Conv2d(256, 256, 3)
+        self.in08_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 6
+        self.refpad09_1 = nn.ReflectionPad2d(1)
+        self.conv09_1 = nn.Conv2d(256, 256, 3)
+        self.in09_1 = InstanceNormalization(256)
+        # relu
+        self.refpad09_2 = nn.ReflectionPad2d(1)
+        self.conv09_2 = nn.Conv2d(256, 256, 3)
+        self.in09_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 7
+        self.refpad10_1 = nn.ReflectionPad2d(1)
+        self.conv10_1 = nn.Conv2d(256, 256, 3)
+        self.in10_1 = InstanceNormalization(256)
+        # relu
+        self.refpad10_2 = nn.ReflectionPad2d(1)
+        self.conv10_2 = nn.Conv2d(256, 256, 3)
+        self.in10_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 8
+        self.refpad11_1 = nn.ReflectionPad2d(1)
+        self.conv11_1 = nn.Conv2d(256, 256, 3)
+        self.in11_1 = InstanceNormalization(256)
+        # relu
+        self.refpad11_2 = nn.ReflectionPad2d(1)
+        self.conv11_2 = nn.Conv2d(256, 256, 3)
+        self.in11_2 = InstanceNormalization(256)
+        # + input
+
+        ##------------------------------------##
+        self.deconv01_1 = nn.ConvTranspose2d(256, 128, 3, 2, 1, 1)
+        self.deconv01_2 = nn.Conv2d(128, 128, 3, 1, 1)
+        self.in12_1 = InstanceNormalization(128)
+        # relu
+        self.deconv02_1 = nn.ConvTranspose2d(128, 64, 3, 2, 1, 1)
+        self.deconv02_2 = nn.Conv2d(64, 64, 3, 1, 1)
+        self.in13_1 = InstanceNormalization(64)
+        # relu
+        self.refpad12_1 = nn.ReflectionPad2d(3)
+        self.deconv03_1 = nn.Conv2d(64, 3, 7)
+        # tanh
+
+    def forward(self, x):
+        y = F.relu(self.in01_1(self.conv01_1(self.refpad01_1(x))))
+        y = F.relu(self.in02_1(self.conv02_2(self.conv02_1(y))))
+        t04 = F.relu(self.in03_1(self.conv03_2(self.conv03_1(y))))
+
+        ##
+        y = F.relu(self.in04_1(self.conv04_1(self.refpad04_1(t04))))
+        t05 = self.in04_2(self.conv04_2(self.refpad04_2(y))) + t04
+
+        y = F.relu(self.in05_1(self.conv05_1(self.refpad05_1(t05))))
+        t06 = self.in05_2(self.conv05_2(self.refpad05_2(y))) + t05
+
+        y = F.relu(self.in06_1(self.conv06_1(self.refpad06_1(t06))))
+        t07 = self.in06_2(self.conv06_2(self.refpad06_2(y))) + t06
+
+        y = F.relu(self.in07_1(self.conv07_1(self.refpad07_1(t07))))
+        t08 = self.in07_2(self.conv07_2(self.refpad07_2(y))) + t07
+
+        y = F.relu(self.in08_1(self.conv08_1(self.refpad08_1(t08))))
+        t09 = self.in08_2(self.conv08_2(self.refpad08_2(y))) + t08
+
+        y = F.relu(self.in09_1(self.conv09_1(self.refpad09_1(t09))))
+        t10 = self.in09_2(self.conv09_2(self.refpad09_2(y))) + t09
+
+        y = F.relu(self.in10_1(self.conv10_1(self.refpad10_1(t10))))
+        t11 = self.in10_2(self.conv10_2(self.refpad10_2(y))) + t10
+
+        y = F.relu(self.in11_1(self.conv11_1(self.refpad11_1(t11))))
+        y = self.in11_2(self.conv11_2(self.refpad11_2(y))) + t11
+        ##
+
+        y = F.relu(self.in12_1(self.deconv01_2(self.deconv01_1(y))))
+        y = F.relu(self.in13_1(self.deconv02_2(self.deconv02_1(y))))
+        y = F.tanh(self.deconv03_1(self.refpad12_1(y)))
+
+        return y
+
+
+class InstanceNormalization(nn.Module):
+    def __init__(self, dim, eps=1e-9):
+        super(InstanceNormalization, self).__init__()
+        self.scale = nn.Parameter(torch.FloatTensor(dim))
+        self.shift = nn.Parameter(torch.FloatTensor(dim))
+        self.eps = eps
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        self.scale.data.uniform_()
+        self.shift.data.zero_()
+
+    def __call__(self, x):
+        n = x.size(2) * x.size(3)
+        t = x.view(x.size(0), x.size(1), n)
+        mean = torch.mean(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x)
+        # Calculate the biased var. torch.var returns unbiased var
+        var = torch.var(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x) * (
+            (n - 1) / float(n)
+        )
+        scale_broadcast = self.scale.unsqueeze(1).unsqueeze(1).unsqueeze(0)
+        scale_broadcast = scale_broadcast.expand_as(x)
+        shift_broadcast = self.shift.unsqueeze(1).unsqueeze(1).unsqueeze(0)
+        shift_broadcast = shift_broadcast.expand_as(x)
+        out = (x - mean) / torch.sqrt(var + self.eps)
+        out = out * scale_broadcast + shift_broadcast
+        return out

requirements.txt

@@ -0,0 +1,3 @@
+torch==1.9.0
+torchvision==0.11.2
+pillow==9.0.0

requirements_dev.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+black==21.12b0
+flake8==4.0.1