first try

Photo_Style_Transfer.ipynb

@@ -0,0 +1,71 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "R7pPsDHPE_PF"
+      },
+      "source": [
+        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dabidou025/Live-Style-Transfer/blob/main/Photo_Style_Transfer.ipynb)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "pRMt1Ae6Asc7",
+        "outputId": "3b3868b0-e6c4-4a6d-93ae-bad6e0e31a02"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install -q gradio==2.8.7\n",
+        "!git clone https://github.com/dabidou025/Live-Style-Transfer.git\n",
+        "%cd Live-Style-Transfer/"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "LBy4brOoF9xC",
+        "outputId": "989467b2-546b-4b80-9475-4ad71a54a937"
+      },
+      "outputs": [],
+      "source": [
+        "from colab_tools_2 import *"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "gradio_pls()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [],
+      "name": "Photo_Style_Transfer.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

colab_tools_2.py

@@ -0,0 +1,87 @@
+# import dependencies
+from IPython.display import display, Javascript, Image
+from google.colab.output import eval_js
+from google.colab.patches import cv2_imshow
+from base64 import b64decode, b64encode
+import cv2
+import numpy as np
+import PIL
+import io
+import html
+import time
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+from models.stmodel import STModel
+from predictor import Predictor
+import argparse
+from glob import glob
+import os
+from ipywidgets import Box, Image
+import gradio as gr
+
+def predict_gradio(image):
+    img_size = 512
+    load_model_path = "./models/st_model_512_80k_12.pth"
+    styles_path = "./styles/" 
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_styles = len(glob(os.path.join(styles_path, '*.jpg')))
+    st_model = STModel(n_styles)
+    if True:
+        st_model.load_state_dict(torch.load(load_model_path, map_location=device))
+    st_model = st_model.to(device)
+
+    predictor = Predictor(st_model, device, img_size)
+
+    list_gen=[]
+    for s in range(n_styles):
+        gen = predictor.eval_image(image, s)
+        list_gen.append(gen)
+    return list_gen
+
+def gradio_pls():
+    description="""
+Upload a photo and click on submit to see the 12 styles applied to your photo. \n 
+Keep in mind that for compatibility reasons your photo is cropped before the neural net applied the different styles.
+<center>
+<table><tr>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/a_muse_picasso.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/britto.jpg" width=100px></td>
+
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/cat.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/cubist.jpg" width=100px></td>
+
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/fractal.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/horse.jpg" width=100px></td>
+
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/monet.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/sketch.jpg" width=100px></td>
+
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/starry_night.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/texture.jpg" width=100px></td>
+
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/tsunami.jpg" width=100px></td>
+<td><img src="https://raw.githubusercontent.com/dabidou025/Live-Style-Transfer/main/styles/vibrant.jpg" width=100px></td>
+
+</tr>
+</table>
+</center>
+"""
+    iface = gr.Interface(
+        predict_gradio,   
+        [
+            gr.inputs.Image(type="pil", label="Image"),
+        ],
+        [
+            gr.outputs.Carousel("image", label="Style"),
+        ],
+        layout="unaligned",
+        title="Photo Style Transfer",
+        description=description,
+        theme="grass",
+        allow_flagging='never'
+    )
+
+    return iface.launch(inline=True, height=800, width=800)

predictor.py

@@ -0,0 +1,51 @@
+import torch
+from torchvision import transforms
+
+from PIL import Image
+
+import numpy as np
+
+class Predictor:
+    def __init__(self, st_model, device, img_size):
+        self.device = device
+
+        self.st_model = st_model.to(device)
+        self.st_model.eval()
+        
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]   
+        
+        self.transformer = transforms.Compose([
+            transforms.Resize(img_size),
+            transforms.CenterCrop(img_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=self.mean, std=self.std)
+        ])
+        
+    def eval_image(self, img, style_1, style_2=None, alpha=0.5):
+        img = self.transformer(img).to(self.device)
+        gen = self.st_model(img.unsqueeze(0), style_1, style_2, alpha)
+        
+        return Image.fromarray(np.uint8(np.moveaxis(gen[0].cpu().detach().numpy()*255.0, 0, 2)))
+
+class WebcamPredictor:
+    def __init__(self, st_model, device):
+        self.device = device
+
+        self.st_model = st_model.to(device)
+        self.st_model.eval()
+        
+        self.mean = np.array([0.485, 0.456, 0.406])
+        self.std = np.array([0.229, 0.224, 0.225])
+
+        self.mean = np.expand_dims(self.mean, (1,2))
+        self.std = np.expand_dims(self.std, (1,2))
+
+    def eval_image(self, img, style_1, style_2=None, alpha=0.5):
+        img = (img - self.mean) / self.std
+        img = torch.from_numpy(img).to(self.device)
+        img = img.float()
+
+        gen = self.st_model(img.unsqueeze(0), style_1, style_2, alpha)
+
+        return np.uint8(gen[0].cpu().detach().numpy()*255.0)

stmodel.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+import math
+import time
+
+class ConvCIN(nn.Module):
+    def __init__(self, n_styles, C_in, C_out, kernel_size, padding, stride, activation=None):
+        super(ConvCIN, self).__init__()
+        
+        self.reflection = nn.ReflectionPad2d(padding)
+        self.conv = nn.Conv2d(in_channels=C_in, out_channels=C_out, kernel_size=kernel_size, stride=stride)
+        nn.init.normal_(self.conv.weight, mean=0, std=1e-2)
+
+        self.instnorm = nn.InstanceNorm2d(C_out)#, affine=True)
+        #nn.init.normal_(self.instnorm.weight, mean=1, std=1e-2)
+        #nn.init.normal_(self.instnorm.bias, mean=0, std=1e-2)
+
+        
+        self.gamma = torch.nn.Parameter(data=torch.randn(n_styles, C_out)*1e-2 + 1, requires_grad=True)
+        #self.gamma.data.uniform_(1.0, 1.0)
+
+        self.beta = torch.nn.Parameter(data=torch.randn(n_styles, C_out)*1e-2, requires_grad=True)
+        #self.beta.data.uniform_(0, 0)
+
+        self.activation = activation
+
+    def forward(self, x, style_1, style_2, alpha):
+
+        x = self.reflection(x)
+        x = self.conv(x)
+
+        x = self.instnorm(x)
+
+        
+        if style_2 != None:
+            gamma = alpha*self.gamma[style_1] + (1-alpha)*self.gamma[style_2]
+            beta = alpha*self.beta[style_1] + (1-alpha)*self.beta[style_2]
+        else:
+            gamma = self.gamma[style_1]
+            beta = self.beta[style_1]
+        
+
+        b,d,w,h = x.size()
+        x = x.view(b,d,w*h)
+
+        x = (x*gamma.unsqueeze(-1) + beta.unsqueeze(-1)).view(b,d,w,h)
+
+        if self.activation == 'relu':
+            x = F.relu(x)
+        elif self.activation == 'sigmoid':
+            x = torch.sigmoid(x)
+
+        return x
+
+class ResidualBlock(nn.Module):
+    def __init__(self, n_styles, C_in, C_out):
+        super(ResidualBlock,self).__init__()
+
+        self.convcin1 = ConvCIN(n_styles, C_in, C_out, kernel_size=3, padding=1, stride=1, activation='relu')
+        self.convcin2 = ConvCIN(n_styles, C_in, C_out, kernel_size=3, padding=1, stride=1)
+
+    def forward(self, x, style_1, style_2, alpha):
+        out = self.convcin1(x, style_1, style_2, alpha)
+        out = self.convcin2(out, style_1, style_2, alpha)
+        return x + out
+
+class UpSampling(nn.Module):
+    def __init__(self, n_styles, C_in, C_out):
+        super(UpSampling,self).__init__()
+
+        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+        self.convcin = ConvCIN(n_styles, C_in, C_out, kernel_size=3, padding=1, stride=1, activation='relu')
+
+    def forward(self, x, style_1, style_2, alpha):
+        x = self.upsample(x)
+        x = self.convcin(x, style_1, style_2, alpha)
+        return x
+
+class STModel(nn.Module):
+    def __init__(self, n_styles):
+        super(STModel,self).__init__()
+
+        self.convcin1 = ConvCIN(n_styles, C_in=3, C_out=32, kernel_size=9, padding=4, stride=1, activation='relu')
+        self.convcin2 = ConvCIN(n_styles, C_in=32, C_out=64, kernel_size=3, padding=1, stride=2, activation='relu')
+        self.convcin3 = ConvCIN(n_styles, C_in=64, C_out=128, kernel_size=3, padding=1, stride=2, activation='relu')
+
+        self.rb1 = ResidualBlock(n_styles, 128, 128)
+        self.rb2 = ResidualBlock(n_styles, 128, 128)
+        self.rb3 = ResidualBlock(n_styles, 128, 128)
+        self.rb4 = ResidualBlock(n_styles, 128, 128)
+        self.rb5 = ResidualBlock(n_styles, 128, 128)
+
+        self.upsample1 = UpSampling(n_styles, 128, 64)
+        self.upsample2 = UpSampling(n_styles, 64, 32)
+
+        self.convcin4 = ConvCIN(n_styles, C_in=32, C_out=3, kernel_size=9, padding=4, stride=1, activation='sigmoid')
+
+    def forward(self, x, style_1, style_2=None, alpha=0.5):
+        x = self.convcin1(x, style_1, style_2, alpha)
+        x = self.convcin2(x, style_1, style_2, alpha)
+        x = self.convcin3(x, style_1, style_2, alpha)
+
+        x = self.rb1(x, style_1, style_2, alpha)
+        x = self.rb2(x, style_1, style_2, alpha)
+        x = self.rb3(x, style_1, style_2, alpha)
+        x = self.rb4(x, style_1, style_2, alpha)
+        x = self.rb5(x, style_1, style_2, alpha)
+
+        x = self.upsample1(x, style_1, style_2, alpha)
+        x = self.upsample2(x, style_1, style_2, alpha)
+
+        x = self.convcin4(x, style_1, style_2, alpha)
+
+        return x