Create model.py

model.py

@@ -0,0 +1,85 @@
+#imports
+import gradio as gr
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from torchvision import transforms
+
+
+#modelling
+class NeuralStyleTransfer(nn.Module):
+  def __init__(self):
+    super(NeuralStyleTransfer, self).__init__()
+    self.model = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1).features
+    # picking the feature layers from the conv layers in the model
+    # this is choosen manually going through all the layers in the model, we can also experiment with this selection  
+    self.feature_layers = [4, 6, 8, 10]
+  
+  def forward(self, x):
+    features = []
+    for layer_num, layer in enumerate(self.model):
+      x = layer(x)
+      #getting the selected layer's output from the model as features
+      if layer_num in self.feature_layers:
+        features.append(x)
+    return features
+
+
+def get_output(style_image, content_image, alpha, beta, step, progress=gr.Progress()):
+  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+  loader = transforms.Compose([transforms.Resize((224,224)),
+                                   transforms.ToTensor()]) #converting to tensor automatically scales the values to 0 and 1
+  style = loader(style_image).to(device)
+  content = loader(content_image).to(device)
+  #starting with content image instead of some random noise image to speed up the process
+  generated = loader(content_image).to(device)
+
+  # setting the generated images values to be tracked and modified while training
+  generated.requires_grad_(True)
+  
+  # densenets weigths need not to be updated
+  model = NeuralStyleTransfer()
+  model.to(device)
+  model.eval()
+  
+  #setting parameters
+  step_count = int(step)
+  learning_rate = 0.001
+
+  #custom loss is defined inside the training loop
+  #the values in the generated matrix needs to be updated by the optimizer 
+  optimizer = torch.optim.Adam([generated], lr = learning_rate) 
+
+  #training
+  for i in progress.tqdm(range(step_count)):
+    style_features = model(style.unsqueeze(0))
+    content_features = model(content.unsqueeze(0))
+    generated_features = model(generated.unsqueeze(0))
+    
+    #content loss
+    content_loss = 0
+    for cf, gf in zip(content_features, generated_features):
+      content_loss += torch.sum((cf-gf)**2)
+
+    #style loss
+    style_loss = 0
+    for sf, gf in zip(style_features, generated_features):
+      bs, c, h, w = sf.shape
+      s_gram = torch.mm(sf.view(c, h*w), sf.view(c, h*w).T)
+      g_gram = torch.mm(gf.view(c, h*w), gf.view(c, h*w).T)
+      style_loss += torch.sum((s_gram - g_gram)**2)
+
+    #total_loss
+    loss = alpha * content_loss + beta * style_loss
+
+    #update values in the generated image
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    
+    if (i+1) % 5 == 0:
+      print(f"\nLoss at {i+1} epoch -----> {loss.item()}", end='')
+  
+  convertor = transforms.ToPILImage() #converts tensor to pil image formate used for displaying in gradio
+  return convertor(generated)