Upload 11 files

app.py

@@ -0,0 +1,70 @@
+
+from model import create_vgg_model
+from data_setup import preprocess, deprocess, get_features, gram_matrix
+from loss_functions import content_loss, style_loss, total_loss
+from torch import optim
+from pathlib import Path
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+
+def predict(content_image, style_image):
+
+    # Create model
+    model = create_vgg_model() 
+
+    # Transform images
+    content_img = preprocess(content_image) 
+    style_img = preprocess(style_image)
+    target_img = content_img.clone().requires_grad_(True)
+
+    content_features = get_features(content_img, model) 
+    style_features = get_features(style_img, model)
+    
+    style_gram = {layer: gram_matrix(style_features[layer]) for layer in style_features}
+
+   # Inference
+    optimizer = optim.Adam([target_img], lr=0.06)
+
+    alpha_param = 1
+    beta_param = 1e2
+    epochs = 60
+  
+    for i in range(epochs):
+      target_features = get_features(target_img, model)
+
+      c_loss = content_loss(target_features['layer_4'], content_features['layer_4'])
+      s_loss = style_loss(target_features, style_gram)
+      t_loss = total_loss(c_loss, s_loss, alpha_param, beta_param)
+
+      optimizer.zero_grad()
+      t_loss.backward()
+      optimizer.step()
+        
+    results = deprocess(target_img)
+
+    return Image.fromarray((results * 255).astype(np.uint8))
+
+# Gradio Interface
+
+example_list =  [['content/content1.jpg',
+                  'style/style1.jpg'],
+                ['content/content2.jpg',
+                  'style/style2.jpg'],
+                ['content/content3.jpg',
+                  'style/style3.jpg']]
+
+title = "Neural Style Transfer 🎨"
+description = "It will take about 1 minute for the result to be displayed. Since the algorithm runs a small number of epochs (to reduce the waiting time), the result will not always be as good as it should be."
+
+demo = gr.Interface(fn=predict,
+                    inputs=['image', 'image'],
+                    outputs=gr.Image().style(width=256, height=256),
+                    examples=example_list,
+                    title=title,
+                    description=description)
+
+
+demo.launch(debug=False,
+            share=False)

data_setup.py

@@ -0,0 +1,61 @@
+
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+
+def preprocess(img):
+
+  image = Image.fromarray(img).convert('RGB')
+  imsize = 196
+  
+  transform = transforms.Compose([
+      transforms.Resize((imsize, imsize)),
+      transforms.ToTensor(),
+      transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                           std=[0.229, 0.224, 0.225])
+  ])
+
+  image = transform(image)
+  image = image.unsqueeze(dim=0)
+
+  return image
+
+
+def deprocess(image): # def show_image
+  
+  image = image.clone()
+  image = image.squeeze(0)
+  image = image.permute(1,2,0)
+  image = image.detach().numpy()
+  image = image * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])
+  image = image.clip(0,1)
+
+  return image
+
+def get_features(image, model):
+
+  features = {}
+  layers = {
+    '0': 'layer_1',
+    '5': 'layer_2',
+    '10': 'layer_3',
+    '19': 'layer_4',
+    '28': 'layer_5'
+    }
+  x = image
+
+  for name, layer in model._modules.items():
+    x = layer(x)
+    if name in layers:
+      features[layers[name]] = x
+
+  return features
+
+
+def gram_matrix(image):
+
+  b, c, h, w = image.size()
+  image = image.view(c, h*w)
+  gram = torch.mm(image, image.t())
+  return gram

loss_functions.py

@@ -0,0 +1,31 @@
+
+import torch
+from data_setup import gram_matrix
+
+def content_loss(target, content):
+
+  loss = torch.mean((target - content) ** 2)
+
+  return loss
+
+
+def style_loss(target_features, style_grams):
+
+  loss = 0
+
+  for layer in target_features:
+    target_f = target_features[layer]
+    target_gram = gram_matrix(target_f)
+    style_gram = style_grams[layer]
+    b,c,h,w = target_f.shape
+    layer_loss = 0.2 * torch.mean((target_gram - style_gram) ** 2)
+    loss += layer_loss/(c*h*w)
+
+  return loss
+
+
+def total_loss(content_loss, style_loss, alpha, beta):
+
+  loss = alpha * content_loss + beta * style_loss
+
+  return loss

model.py

@@ -0,0 +1,18 @@
+
+import torchvision
+
+def create_vgg_model():
+   
+    # Create model
+    model_weights = torchvision.models.VGG19_Weights.DEFAULT
+    model = torchvision.models.vgg19(weights=model_weights)
+    
+
+    # Freeze layers
+    for param in model.parameters():
+      param.requires_grad = False
+
+    # Kepp only the features of the model
+    model = model.features
+    
+    return model

requirements.txt

@@ -0,0 +1,2 @@
+torch==1.12.0
+torchvision==0.13.0gradio==3.1.4