Upload 2 files

app.py

@@ -0,0 +1,140 @@
+from __future__ import print_function
+import torch
+import torchvision
+import torch.nn as nn
+import gradio as gr
+import os
+import time
+import torch.nn.functional as F
+import torch.optim as optim
+import matplotlib.pyplot as plt
+import torchvision.transforms as transforms
+import copy
+import torchvision.models as models
+import torchvision.transforms.functional as TF
+from PIL import Image
+import numpy as np
+from model import ContentLoss, gram_matrix, StyleLoss, image_transform, get_input_optimizer,get_style_model_and_losses
+
+
+
+
+
+## style_transfer
+import numpy as np
+def run_style_transfer(cnn, normalization_mean, normalization_std,
+                       content_img, style_img, input_img, num_steps=300,
+                       style_weight=1000000, content_weight=1):
+    """Run the style transfer."""
+    print('Building the style transfer model..')
+    model, style_losses, content_losses = get_style_model_and_losses(cnn,
+        normalization_mean, normalization_std, style_img, content_img)
+
+    # We want to optimize the input and not the model parameters so we
+    # update all the requires_grad fields accordingly
+    input_img.requires_grad_(True)
+    model.requires_grad_(False)
+
+    optimizer = get_input_optimizer(input_img)
+
+    print('Optimizing..')
+    run = [0]
+    while run[0] <= num_steps:
+
+        def closure():
+            # correct the values of updated input image
+            with torch.no_grad():
+                input_img.clamp_(0, 1)
+
+            optimizer.zero_grad()
+            model(input_img)
+            style_score = 0
+            content_score = 0
+
+            for sl in style_losses:
+                style_score += sl.loss
+            for cl in content_losses:
+                content_score += cl.loss
+
+            style_score *= style_weight
+            content_score *= content_weight
+
+            loss = style_score + content_score
+            loss.backward()
+
+            run[0] += 1
+            if run[0] % 50 == 0:
+                print("run {}:".format(run))
+                print('Style Loss : {:4f} Content Loss: {:4f}'.format(
+                    style_score.item(), content_score.item()))
+                print()
+
+            return style_score + content_score
+
+        optimizer.step(closure)
+
+    # a last correction...
+    with torch.no_grad():
+        input_img.clamp_(0, 1)
+
+    # Convert output tensor to a NumPy array
+    output_np = input_img.detach().cpu().numpy()[0].transpose(1, 2, 0)
+
+    # Convert NumPy array to PIL Image object
+    output_img = Image.fromarray((output_np * 255).astype(np.uint8))
+    
+    return output_img
+        
+             
+#Defining the predict function
+def style_transfer(cont_img,styl_img):
+  
+  #Start the timer
+  start_time = time.time()
+  
+
+  #transform the input image
+  style_img = image_transform(styl_img)
+  content_img =image_transform(cont_img)
+  
+  #getting input image
+  input_img = content_img.clone()
+
+  #running the style transfer
+  output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std,
+                            content_img, style_img, input_img)
+  # output_img = output.detach().cpu().squeeze(0)
+  # output_img = TF.to_pil_image(output_img)
+  end_time=time.time()
+  
+  pred_time =round(end_time- start_time, 5)
+  
+  return output
+
+##Gradio App
+import gradio as gr
+title= 'Style Transfer'
+description='A model to transfer the style of one image to another'
+article = 'Created at Pytorch Model Deployment'
+
+#example_images
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+#Create the gradio demo
+demo = gr.Interface(
+    fn=style_transfer,
+    inputs=[
+        gr.inputs.Image(label="content Image",type='pil'),
+        gr.inputs.Image(label="style_image",type='pil')
+    ],
+    examples=example_list,
+    outputs="image",
+    allow_flagging=False,
+    title=title,
+    description=description,
+    article=article
+)
+
+# Launch the Gradio interface
+demo.launch(debug=True)
+

model.py

@@ -0,0 +1,135 @@
+
+#Content Loss
+class ContentLoss(nn.Module):
+
+    def __init__(self, target,):
+        super(ContentLoss, self).__init__()
+        '''
+        we 'detach' the target content from the tree used 
+        to dynamically compute the gradient: this is a stated value,
+        not a variable. Otherwise the forward method of the criterion
+        will throw an error.
+        '''
+        self.target = target.detach()
+      
+    def forward(self, input):
+        self.loss = F.mse_loss(input, self.target)
+        return input
+
+#Style Loss
+def gram_matrix(input):
+    a, b, c, d = input.size()  # a=batch size(=1)
+    # b=number of feature maps
+    # (c,d)=dimensions of a f. map (N=c*d)
+
+    features = input.view(a * b, c * d)  # resize F_XL into \hat F_XL
+
+    G = torch.mm(features, features.t())  # compute the gram product
+
+    # we 'normalize' the values of the gram matrix
+    # by dividing by the number of element in each feature maps.
+    return G.div(a * b * c * d)
+
+class StyleLoss(nn.Module):
+
+    def __init__(self, target_feature):
+        super(StyleLoss, self).__init__()
+        self.target = gram_matrix(target_feature).detach()
+
+    def forward(self, input):
+        G = gram_matrix(input)
+        self.loss = F.mse_loss(G, self.target)
+        return input  
+
+
+#Normalization
+
+cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
+cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
+#image transformation
+def image_transform(image):
+    if isinstance(image, str):
+        # If image is a path to a file, open it using PIL
+        image = Image.open(image).convert('RGB')
+    else:
+        # If image is a NumPy array, convert it to a PIL image
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+    # Apply the same transformations as before
+    image = transform(image).unsqueeze(0)
+    return image.to(device)
+
+
+
+#Defining a model
+cnn = models.vgg19(pretrained=True).features.to(device).eval()
+
+#getting the input optimizer
+def get_input_optimizer(input_img):
+    # this line to show that input is a parameter that requires a gradient
+    optimizer = optim.LBFGS([input_img])
+    return optimizer
+
+# desired depth layers to compute style/content losses :
+
+
+content_layers_default = ['conv_4']
+style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
+
+def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
+                               style_img, content_img,
+                               content_layers=content_layers_default,
+                               style_layers=style_layers_default):
+    # normalization module
+    normalization = Normalization(normalization_mean, normalization_std).to(device)
+
+    # just in order to have an iterable access to or list of content/style
+    # losses
+    content_losses = []
+    style_losses = []
+
+    # assuming that ``cnn`` is a ``nn.Sequential``, so we make a new ``nn.Sequential``
+    # to put in modules that are supposed to be activated sequentially
+    model = nn.Sequential(normalization)
+
+    i = 0  # increment every time we see a conv
+    for layer in cnn.children():
+        if isinstance(layer, nn.Conv2d):
+            i += 1
+            name = 'conv_{}'.format(i)
+        elif isinstance(layer, nn.ReLU):
+            name = 'relu_{}'.format(i)
+            # The in-place version doesn't play very nicely with the ``ContentLoss``
+            # and ``StyleLoss`` we insert below. So we replace with out-of-place
+            # ones here.
+            layer = nn.ReLU(inplace=False)
+        elif isinstance(layer, nn.MaxPool2d):
+            name = 'pool_{}'.format(i)
+        elif isinstance(layer, nn.BatchNorm2d):
+            name = 'bn_{}'.format(i)
+        else:
+            raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))
+
+        model.add_module(name, layer)
+
+        if name in content_layers:
+            # add content loss:
+            target = model(content_img).detach()
+            content_loss = ContentLoss(target)
+            model.add_module("content_loss_{}".format(i), content_loss)
+            content_losses.append(content_loss)
+
+        if name in style_layers:
+            # add style loss:
+            target_feature = model(style_img).detach()
+            style_loss = StyleLoss(target_feature)
+            model.add_module("style_loss_{}".format(i), style_loss)
+            style_losses.append(style_loss)
+
+    # now we trim off the layers after the last content and style losses
+    for i in range(len(model) - 1, -1, -1):
+        if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):
+            break
+
+    model = model[:(i + 1)]
+
+    return model, style_losses, content_losses