Initial commit

.gitignore

@@ -0,0 +1,3 @@
+NeuralStyleTransfer.ipynb
+images
+__pycache__

app.py

@@ -0,0 +1,101 @@
+import gradio as gr
+from model import NeuralStyleTransfer
+import tensorflow as tf
+
+
+def model_fn(
+    style, content, extractor="inception_v3", n_content_layers=3, n_style_layers=2,
+    epochs=4, learning_rate=60.0, steps_per_epoch=100, style_weight=1e-2,
+):
+    model = NeuralStyleTransfer(
+        style_image=style,
+        content_image=content,
+        extractor=extractor,
+        n_content_layers=n_content_layers,
+        n_style_layers=n_style_layers,
+    )
+
+    return model.fit_style_transfer(
+        epochs=10,
+        learning_rate=80.0,
+        steps_per_epoch=100,
+        style_weight=1e-2,
+        content_weight=1e-4,
+        show_image=True,
+        show_interval=90,
+        var_weight=1e-12,
+        terminal=False,
+    )
+
+
+def hugging_face():
+    demo = gr.Interface(
+        fn=model_fn,
+        inputs=[
+            "image",
+            "image",
+            gr.Dropdown(
+                ["inception_v3", "vgg19", "resnet50", "mobilenet_v2"],
+                label="extractor",
+                default="inception_v3",
+                info="Feature extractor to use.",
+            ),
+            gr.Slider(
+                1,
+                5,
+                value=3,
+                label="n_content_layers",
+                info="Number of content layers to use.",
+            ),
+            gr.Slider(
+                1,
+                5,
+                value=2,
+                label="n_style_layers",
+                info="Number of style layers to use.",
+            ),
+            gr.Slider(
+                2, 20, value=4, label="epochs", info="Number of epochs to train for."
+            ),
+            gr.Slider(
+                1, 100, value=60, label="learning_rate", info="Initial Learning rate."
+            ),
+            gr.Slider(
+                1,
+                100,
+                value=100,
+                label="steps_per_epoch",
+                info="Number of steps per epoch.",
+            ),
+            gr.Slider(
+                1e-4,
+                1e-2,
+                value=1e-2,
+                label="style_weight",
+                info="Weight of style loss.",
+            ),
+            gr.Slider(
+                1e-4,
+                1e-2,
+                value=1e-4,
+                label="content_weight",
+                info="Weight of content loss.",
+            ),
+            gr.Slider(
+                1e-12,
+                1e-9,
+                value=1e-12,
+                label="var_weight",
+                info="Weight of total variation loss.",
+            ),
+        ],
+        outputs="image",
+    )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = hugging_face()
+    demo.launch( share=True)
+    

flagged/log.csv

@@ -0,0 +1,2 @@
+style,content,extractor,n_content_layers,n_style_layers,epochs,learning_rate,steps_per_epoch,style_weight,content_weight,var_weight,output,flag,username,timestamp
+/home/shailja/Courses/Notes/Projects/StyleTransfer/flagged/style/tmpesg98402.jpg,/home/shailja/Courses/Notes/Projects/StyleTransfer/flagged/content/tmpu5ej5fhy.jpg,inception_v3,3,2,4,60,100,0.01,0.00039,5.8e-11,,,,2023-05-28 08:18:26.193609

model.py

@@ -0,0 +1,354 @@
+import tensorflow as tf
+import imageio
+import numpy as np
+import matplotlib.pyplot as plt
+from keras import backend as K
+from IPython.display import display as display_fn
+
+
+class NeuralStyleTransfer:
+    def __init__(self, style_image, content_image, extractor, n_style_layers=5, n_content_layers=5, display=True):
+        # load the model
+        if extractor == "inception_v3":
+            self.feature_extractor = tf.keras.applications.InceptionV3(
+                include_top=False, weights="imagenet"
+            )
+        elif isinstance(extractor, tf.keras.Model):
+            self.feature_extractor = extractor
+        else:
+            raise Exception("Features Extractor not found")
+
+        # freeze the model
+        self.feature_extractor.trainable = False
+
+        # define the style and content depth
+        self.n_style_layers = n_style_layers
+        self.n_content_layers = n_content_layers
+
+        self.style_image = self._load_img(style_image)
+        self.content_image = self._load_img(content_image)
+
+
+        if display:
+            self.show_images_with_objects(
+                [self.style_image, self.content_image],
+                ["Style Image", "Content Image"],
+            )
+
+    def tensor_to_image(self, tensor):
+        """converts a tensor to an image"""
+        tensor_shape = tf.shape(tensor)
+        number_elem_shape = tf.shape(tensor_shape)
+        if number_elem_shape > 3:
+            assert tensor_shape[0] == 1
+            tensor = tensor[0]
+        return tf.keras.preprocessing.image.array_to_img(tensor)
+
+    def _load_img(self, image):
+        max_dim = 512
+
+        image = tf.io.read_file(image)
+        image = tf.image.decode_image(image)
+        image = tf.image.convert_image_dtype(image, tf.float32)
+
+        image = tf.image.convert_image_dtype(image, tf.float32)
+
+        shape = tf.shape(image)[:-1]
+        shape = tf.cast(tf.shape(image)[:-1], tf.float32)
+        long_dim = max(shape)
+        scale = max_dim / long_dim
+
+        new_shape = tf.cast(shape * scale, tf.int32)
+
+        image = tf.image.resize(image, new_shape)
+        image = image[tf.newaxis, :]
+        image = tf.image.convert_image_dtype(image, tf.uint8)
+
+        return image
+
+    def imshow(self, image, title=None):
+        """displays an image with a corresponding title"""
+        if len(image.shape) > 3:
+            image = tf.squeeze(image, axis=0)
+
+        plt.imshow(image)
+        if title:
+            plt.title(title)
+
+    def show_images_with_objects(self, images, titles=[]):
+        """displays a row of images with corresponding titles"""
+        if len(images) != len(titles):
+            return
+
+        plt.figure(figsize=(20, 12))
+        for idx, (image, title) in enumerate(zip(images, titles)):
+            plt.subplot(1, len(images), idx + 1)
+            plt.xticks([])
+            plt.yticks([])
+            self.imshow(image, title)
+
+    def _preprocess_image(self, image):
+        image = tf.cast(image, dtype=tf.float32)
+        image = (image / 127.5) - 1.0
+
+        return image
+
+    def get_output_layers(self):
+        # get all the layers which contain conv in their name
+        all_layers = [
+            layer.name
+            for layer in self.feature_extractor.layers
+            if "conv" in layer.name
+        ]
+
+        # define the style layers
+        style_layers = all_layers[: self.n_style_layers]
+
+        # define the content layers from second last layer
+        content_layers = all_layers[-2: -self.n_content_layers - 2 : -1]
+
+        content_and_style_layers = content_layers + style_layers
+
+        return content_and_style_layers
+
+    def build(self, layers_name):
+
+        output_layers = [
+            self.feature_extractor.get_layer(name).output for name in layers_name
+        ]
+
+        model = tf.keras.Model(self.feature_extractor.input, output_layers)
+
+        self.feature_extractor = model
+
+        return
+
+    def _loss(self, target_img, features_img, type):
+        """
+        Calculates the loss of the style transfer
+
+        target_img:
+            the target image (style or content) features
+
+        features_img:
+            the generated image features (style or content)
+
+        """
+
+        loss = tf.reduce_mean(tf.square(features_img - target_img))
+
+        if type == "content":
+            return 0.5 * loss
+
+        return loss
+
+    def _gram_matrix(self, input_tensor):
+        """
+        Calculates the gram matrix and divides by the number of locations
+
+        input_tensor:
+            the output of the conv layer of the style image, shape = (batch_size, height, width, channels)
+
+        """
+        result = tf.linalg.einsum("bijc,bijd->bcd", input_tensor, input_tensor)
+        input_shape = tf.shape(input_tensor)
+        num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
+        return result / (num_locations)
+
+    def get_features(self, image, type):
+        preprocess_image = self._preprocess_image(image)
+
+        outputs = self.feature_extractor(preprocess_image)
+
+        if type == "style":
+            outputs = outputs[self.n_content_layers : ]
+            features = [self._gram_matrix(style_output) for style_output in outputs]
+
+        elif type == "content":
+            features = outputs[ : self.n_content_layers]
+
+        return features
+
+    def _style_content_loss(
+        self,
+        style_targets,
+        style_outputs,
+        content_targets,
+        content_outputs,
+        style_weight,
+        content_weight,
+    ):
+        """
+        Calculates the total loss of the style transfer
+
+        style_targets:
+            the style features of the style image
+
+        style_outputs:
+            the style features of the generated image
+
+        content_targets:
+            the content features of the content image
+
+        content_outputs:
+            the content features of the generated image
+
+        style_weight:
+            the weight of the style loss
+
+        content_weight:
+            the weight of the content loss
+
+        """
+
+        # adding the loss of each layer
+        style_loss = style_weight * tf.add_n(
+            [
+                self._loss(style_target, style_output, type="style")
+                for style_target, style_output in zip(style_targets, style_outputs)
+            ]
+        )
+        content_loss = content_weight * tf.add_n(
+            [
+                self._loss(content_target, content_output, type="content")
+                for content_target, content_output in zip(
+                    content_targets, content_outputs
+                )
+            ]
+        )
+        total_loss = style_loss + content_loss
+        return total_loss
+
+    def _grad_loss(
+        self,
+        generated_image,
+        style_target,
+        content_target,
+        style_weight,
+        content_weight,
+        var_weight,
+    ):
+        """
+        Calculates the gradients of the loss function with respect to the generated image
+
+        generated_image:
+            the generated image
+
+        """
+
+        with tf.GradientTape() as tape:
+            style_features = self.get_features(generated_image, type="style")
+            content_features = self.get_features(generated_image, type="content")
+            loss = self._style_content_loss(
+                style_target,
+                style_features,
+                content_target,
+                content_features,
+                style_weight,
+                content_weight,
+            )
+
+            loss += var_weight*tf.image.total_variation(generated_image)
+        grads = tape.gradient(loss, generated_image)
+        return grads, loss
+
+    def _update_image_with_style(
+        self,
+        generated_image,
+        style_target,
+        content_target,
+        style_weight,
+        content_weight,
+        optimizer,
+        var_weight,
+    ):
+        grads, loss = self._grad_loss(
+            generated_image, style_target, content_target, style_weight, content_weight, var_weight
+        )
+
+        optimizer.apply_gradients([(grads, generated_image)])
+
+        generated_image.assign(
+            tf.clip_by_value(generated_image, clip_value_min=0.0, clip_value_max=255.0)
+        )
+        return loss
+
+    def fit_style_transfer(
+        self,
+        epochs=10,
+        learning_rate=80,
+        steps_per_epoch=100,
+        style_weight=1e-2,
+        content_weight=1e-4,
+        show_interval=10,
+        var_weight=0.0,
+    ):
+        """
+        epochs:
+            the number of epochs to train the model for
+
+        learning_rate:
+            the initial learning rate of the optimizer  (default: 80)
+
+        steps_per_epoch:
+            the number of steps to train the model for per epoch
+
+        style_weight:
+            the weight of the style loss
+
+        content_weight:
+            the weight of the content loss
+
+        show_image:
+            whether to save the generated image after each epoch
+
+        show_interval:
+            the interval at which to save the generated image
+
+        var_weight:
+            the weight of the total variation loss
+
+        """
+
+        style_image = self.style_image
+        content_image = self.content_image
+
+        content_and_style_layers = self.get_output_layers()
+
+        # build the model with the layers we need to extract the features from
+        K.clear_session()
+        self.build(content_and_style_layers)
+
+        style_features = self.get_features(style_image, type="style")
+        content_features = self.get_features(content_image, type="content")
+
+        optimizer = tf.optimizers.Adam(
+            tf.keras.optimizers.schedules.ExponentialDecay(
+                initial_learning_rate=learning_rate, decay_steps=100, decay_rate=0.80
+            )
+        )
+
+        generated_image = tf.cast(content_image, tf.float32)
+        generated_image = tf.Variable(generated_image)
+
+        step = 0
+        images = []
+
+        img = None
+
+        for epoch in range(epochs):
+            for step in range(steps_per_epoch):
+                loss = self._update_image_with_style(
+                    generated_image,
+                    style_features,
+                    content_features,
+                    style_weight,
+                    content_weight,
+                    optimizer,
+                    var_weight,
+                )
+
+                display_image = self.tensor_to_image(generated_image)
+
+
+    

requirements.txt

@@ -0,0 +1,2 @@
+tensorflow-cpu
+gradio