adding files

README.md

@@ -1,8 +1,8 @@
 ---
 title: Fast Neural Style Transfer
-emoji: 📊
-colorFrom: red
-colorTo: pink
+emoji: 🎨
+colorFrom: green
+colorTo: red
 sdk: gradio
 app_file: app.py
 pinned: false

app.py

@@ -0,0 +1,79 @@
+from pydoc import describe
+import re
+import numpy as np
+from PIL import Image
+import torch
+from torchvision import transforms
+import gradio as gr
+from model import TransformerNet
+
+style_model = TransformerNet()
+device=torch.device("cpu")
+
+styles_map = {"Kandinsky, Several circles": "kand_circles.model", 
+              "Haring, Dance": "haring_dance.model", 
+              "Picasso, The weeping woman": "picasso_weeping.model",
+              "Van Gogh, Wheatfield with crows": "vangogh_crows.model"}
+
+content_transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Lambda(lambda x: x.mul(255))
+])
+
+def run(content_image, style):
+    content_image.thumbnail((1080, 1080))
+    img = content_transform(content_image)
+    img = img.unsqueeze(0).to(device)
+
+    model = styles_map[style]
+    state_dict = torch.load(f"./models/{model}")
+    
+    for k in list(state_dict.keys()):
+        if re.search(r'in\d+\.running_(mean|var)$', k):
+            del state_dict[k]
+    style_model.load_state_dict(state_dict)
+    style_model.to(device)
+
+    with torch.no_grad():
+        output = style_model(img)
+
+    img = output[0].clone().clamp(0, 255).numpy()
+    img = img.transpose(1, 2, 0).astype("uint8")
+    
+    img = Image.fromarray(img)
+    return img
+
+content_image_input = gr.inputs.Image(label="Content Image", type="pil")
+style_input = gr.inputs.Dropdown(list(styles_map.keys()), type="value", default="Kandinsky, Several circles", label="Style")
+
+description="Fast Neural Style Transfer demo (trained from scratch!). Upload a content image. Select an artwork. Enjoy."
+article="""
+**References**\n\n
+You can find <a href='https://francescopochetti.com/fast-neural-style-transfer-deploying-pytorch-models-to-aws-lambda/' target='_blank'>here</a> a post I put together 
+describing the approach I used to train models and deploy them on <a href='http://visualneurons.com/fast.html' target='_blank'>visualneurons.com</a> using AWS Lambda. \n
+<a href='https://github.com/FraPochetti/examples/blob/master/fast_neural_style/neural_style/FastStyleTransferPytorch.ipynb' target='_blank'>Here</a> is instead the Jupyter notebook
+with the training logic. \n
+<br>
+<hr>
+**Kandinsky, Several circles**
+<img src='https://style-transfer-webapptest.s3.eu-west-1.amazonaws.com/small_images_hf/Several_Circles.jpeg'>
+<hr>
+**Haring, Dance**
+<img src='https://style-transfer-webapptest.s3.eu-west-1.amazonaws.com/small_images_hf/Haring.jpeg'>
+<hr>
+**Picasso, The weeping woman**
+<img src='https://style-transfer-webapptest.s3.eu-west-1.amazonaws.com/small_images_hf/weeping.png'>
+<hr>
+**Van Gogh, Wheatfield with crows**
+<img src='https://style-transfer-webapptest.s3.eu-west-1.amazonaws.com/small_images_hf/Wheatfield_with_crows.jpeg'>
+"""
+example = ["dog.jpeg", "Kandinsky, Several circles"]
+
+app_interface = gr.Interface(fn=run,
+                             inputs=[content_image_input, style_input],
+                             outputs="image",
+                             title="Fast Neural Style Transfer",
+                             description=description,
+                             examples=[example],
+                             article=article)
+app_interface.launch()

model.py

@@ -0,0 +1,98 @@
+import torch
+
+class TransformerNet(torch.nn.Module):
+    def __init__(self):
+        super(TransformerNet, self).__init__()
+        # Initial convolution layers
+        self.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1)
+        self.in1 = torch.nn.InstanceNorm2d(32, affine=True)
+        self.conv2 = ConvLayer(32, 64, kernel_size=3, stride=2)
+        self.in2 = torch.nn.InstanceNorm2d(64, affine=True)
+        self.conv3 = ConvLayer(64, 128, kernel_size=3, stride=2)
+        self.in3 = torch.nn.InstanceNorm2d(128, affine=True)
+        # Residual layers
+        self.res1 = ResidualBlock(128)
+        self.res2 = ResidualBlock(128)
+        self.res3 = ResidualBlock(128)
+        self.res4 = ResidualBlock(128)
+        self.res5 = ResidualBlock(128)
+        # Upsampling Layers
+        self.deconv1 = UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2)
+        self.in4 = torch.nn.InstanceNorm2d(64, affine=True)
+        self.deconv2 = UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2)
+        self.in5 = torch.nn.InstanceNorm2d(32, affine=True)
+        self.deconv3 = ConvLayer(32, 3, kernel_size=9, stride=1)
+        # Non-linearities
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, X):
+        y = self.relu(self.in1(self.conv1(X)))
+        y = self.relu(self.in2(self.conv2(y)))
+        y = self.relu(self.in3(self.conv3(y)))
+        y = self.res1(y)
+        y = self.res2(y)
+        y = self.res3(y)
+        y = self.res4(y)
+        y = self.res5(y)
+        y = self.relu(self.in4(self.deconv1(y)))
+        y = self.relu(self.in5(self.deconv2(y)))
+        y = self.deconv3(y)
+        return y
+
+
+class ConvLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(ConvLayer, self).__init__()
+        reflection_padding = kernel_size // 2
+        self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)
+        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)
+
+    def forward(self, x):
+        out = self.reflection_pad(x)
+        out = self.conv2d(out)
+        return out
+
+
+class ResidualBlock(torch.nn.Module):
+    """ResidualBlock
+    introduced in: https://arxiv.org/abs/1512.03385
+    recommended architecture: http://torch.ch/blog/2016/02/04/resnets.html
+    """
+
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)
+        self.in1 = torch.nn.InstanceNorm2d(channels, affine=True)
+        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
+        self.in2 = torch.nn.InstanceNorm2d(channels, affine=True)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        residual = x
+        out = self.relu(self.in1(self.conv1(x)))
+        out = self.in2(self.conv2(out))
+        out = out + residual
+        return out
+
+
+class UpsampleConvLayer(torch.nn.Module):
+    """UpsampleConvLayer
+    Upsamples the input and then does a convolution. This method gives better results
+    compared to ConvTranspose2d.
+    ref: http://distill.pub/2016/deconv-checkerboard/
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
+        super(UpsampleConvLayer, self).__init__()
+        self.upsample = upsample
+        reflection_padding = kernel_size // 2
+        self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)
+        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)
+
+    def forward(self, x):
+        x_in = x
+        if self.upsample:
+            x_in = torch.nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)
+        out = self.reflection_pad(x_in)
+        out = self.conv2d(out)
+        return out

requirements.txt

@@ -0,0 +1,3 @@
+pillow > 8.0.0,<9
+torchvision==0.10.0
+torch===1.9.0