Added files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,176 @@
+import gradio as gr
+from style_transfer import StyleTransfer
+import tensorflow as tf
+from tensorflow.keras import backend as K
+import numpy as np
+
+
+def validate_inputs(epochs, steps_per_epoch, image_frequency, alpha, beta, lr):
+    """Validates the inputs and converts them to the correct type"""
+    epochs = int(epochs)
+    steps_per_epoch = int(steps_per_epoch)
+    image_frequency = int(image_frequency)
+    alpha = float(alpha)
+    beta = float(beta)
+    lr = float(lr)
+    return epochs, steps_per_epoch, image_frequency, alpha, beta, lr
+
+
+def stylize_image(
+    content_image_path,
+    style_image_path,
+    epochs,
+    steps_per_epoch,
+    image_frequency,
+    alpha,
+    beta,
+    lr,
+):
+    """Stylizes the image using the style and content images
+
+    Parameters
+    ----------
+    content_image_path : str
+        Path to the content image
+    style_image_path : str
+        Path to the style image
+    epochs : int, optional
+        Number of epochs
+    steps_per_epoch : int, optional
+        Number of steps per epoch
+    image_frequency : int, optional
+        Frequency of images to show
+    alpha : float, optional
+        Content weight
+    beta : float, optional
+        Style weight
+    lr : float, optional
+        Learning rate
+
+    Returns
+    -------
+    [PIL.Image]
+        List of images
+    """
+    epochs, steps_per_epoch, image_frequency, alpha, beta, lr = validate_inputs(
+        epochs, steps_per_epoch, image_frequency, alpha, beta, lr
+    )
+    style_transfer = StyleTransfer(
+        content_image_path=content_image_path,
+        style_image_path=style_image_path,
+    )
+    if style_transfer.model is None:
+        K.clear_session()
+        _ = style_transfer.load_model()
+    style_image = style_transfer.load_image(style_transfer.style_image_path)
+    content_image = style_transfer.load_image(style_transfer.content_image_path)
+
+    style_target = style_transfer.get_features(style_image, "style")
+    content_target = style_transfer.get_features(content_image, "content")
+
+    target = content_target + style_target
+    image = tf.cast(content_image, dtype=tf.float32)
+    image = tf.Variable(image)
+    optimizer = tf.optimizers.Adam(
+        tf.keras.optimizers.schedules.ExponentialDecay(
+            initial_learning_rate=lr, decay_steps=100, decay_rate=0.80
+        )
+    )
+    for epoch in range(epochs):
+        for step in range(steps_per_epoch):
+            loss = style_transfer.update_image(image, target, alpha, beta, optimizer)
+            display_image = style_transfer.tensor_to_image(image)
+            # images.append(display_image)
+            if (step) % image_frequency == 0:
+                yield np.array(display_image), epoch + 1, step + 1, loss
+
+
+def main():
+    content_image = gr.Image(type="filepath", label="Content Image", shape=(512, 512))
+    style_image = gr.Image(type="filepath", label="Style Image", shape=(512, 512))
+    epochs = gr.Slider(minimum=1, maximum=20, label="Epochs", value=10)
+    steps_per_epoch = gr.Slider(
+        minimum=1, maximum=20, label="Steps per Epoch", value=10
+    )
+    image_frequency = gr.Slider(
+        minimum=1, maximum=10, label="Show Image Frequency", value=2
+    )
+    alpha = gr.Slider(minimum=0, maximum=1, label="Alpha", value=1)
+    beta = gr.Slider(minimum=0, maximum=1, label="Beta", value=0.1)
+    lr = gr.Slider(minimum=0.1, maximum=100, label="Learning Rate", value=40.0)
+
+    output_image = gr.Image(type="numpy", label="Output Image", shape=(512, 512))
+    current_epoch = gr.Number(label="Current Epoch")
+    current_step = gr.Number(label="Current Step")
+    current_loss = gr.Number(label="Current Loss")
+
+    inputs = [
+        content_image,
+        style_image,
+        epochs,
+        steps_per_epoch,
+        image_frequency,
+        alpha,
+        beta,
+        lr,
+    ]
+
+    outputs = [output_image, current_epoch, current_step, current_loss]
+
+    description = """### This is a demo of neural style transfer. Upload a content image and a style image, and see the result! You can play around with the parameters to see how they affect the result.
+    """
+
+    interface = gr.Interface(
+        fn=stylize_image,
+        inputs=inputs,
+        outputs=outputs,
+        title="Style Transfer",
+        description=description,
+        examples=[
+            [
+                "examples/landscape_1.jpg",
+                "examples/van_gogh.jpg",
+                10,
+                10,
+                1,
+                1,
+                0.1,
+                30.0,
+            ],
+            [
+                "examples/landscape_1.jpg",
+                "examples/picaso.jpg",
+                10,
+                10,
+                1,
+                1,
+                0.1,
+                30.0,
+            ],
+            [
+                "examples/landscape_2.jpg",
+                "examples/van_gogh.jpg",
+                10,
+                10,
+                1,
+                1,
+                0.1,
+                30.0,
+            ],
+            [
+                "examples/landscape_2.jpg",
+                "examples/picaso.jpg",
+                10,
+                10,
+                1,
+                1,
+                0.1,
+                30.0,
+            ],
+        ],
+        theme="gstaff/xkcd",
+    )
+    interface.queue().launch(server_name="0.0.0.0", server_port=7860)
+
+
+main()

requirements.txt

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

style_transfer.py

@@ -0,0 +1,330 @@
+import tensorflow as tf
+import numpy as np
+import matplotlib.pyplot as plt
+from keras import backend as K
+
+
+class StyleTransfer:
+    """A class for neural style transfer. Uses the Inception model to extract features from the content and style images."""
+
+    content_layers = ["conv2d_88", "conv2d_91", "conv2d_92", "conv2d_85", "conv2d_93"]
+    style_layers = ["conv2d", "conv2d_1", "conv2d_2", "conv2d_3", "conv2d_4"]
+    content_and_style_layers = content_layers + style_layers
+
+    NUM_CONTENT_LAYERS = len(content_layers)
+    NUM_STYLE_LAYERS = len(style_layers)
+
+    def __init__(self, content_image_path, style_image_path) -> None:
+        """Initializes the class
+
+        Parameters
+        ----------
+        content_image_path : str
+            path to the content image
+        style_image_path : str
+            path to the style image
+
+        Returns
+        -------
+        None
+        """
+        self.content_image_path = content_image_path
+        self.style_image_path = style_image_path
+        self.model = None
+
+    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, "There are more than one image"
+            tensor = tensor[0]
+        return tf.keras.preprocessing.image.array_to_img(tensor)
+
+    def load_image(self, path_to_img):
+        """loads an image as a tensor and scales it to 512 pixels"""
+        max_dim = 512
+        image = tf.io.read_file(path_to_img)
+        image = tf.image.decode_jpeg(image)
+        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=""):
+        """displays an image"""
+        if len(image.shape) > 3:
+            image = tf.squeeze(image, axis=0)
+
+        plt.imshow(image)
+        plt.title(title)
+
+    def show_images_with_style(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)
+        plt.show()
+
+    def preprocess_image(self, image):
+        """preprocesses a given image to use with Inception model"""
+        image = tf.cast(image, dtype=tf.float32)
+        image = (image / 127.5) - 1.0
+
+        return image
+
+    def display_images(self):
+        """displays the content and style images"""
+        content_image = self.load_image(self.content_image_path)
+        style_image = self.load_image(self.style_image_path)
+
+        self.show_images_with_style(
+            [content_image, style_image],
+            titles=[f"Content image", f"Style image"],
+        )
+
+    def gram_matrix(self, input_tensor):
+        """Calculates the gram matrix and divides by the number of locations
+
+        Parameters
+        ----------
+        input_tensor : tensor
+            tensor to calculate the gram matrix from
+
+        Returns
+        -------
+        tensor
+            gram matrix of the input tensor
+        """
+
+        gram = 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)
+        scaled_gram = gram / num_locations
+
+        return scaled_gram
+
+    def get_features(self, image, type=None):
+        """Returns the features of the image
+
+        Parameters
+        ----------
+        image : tensor
+            image to extract features from
+        type : str
+            type of features to extract. Either "style" or "content". If `None` is provided, both
+            content and style features are returned
+
+        Returns
+        -------
+        list
+            list of features of the content and style images
+        """
+        preprocessed_image = self.preprocess_image(image)
+        outputs = self.model(preprocessed_image)
+
+        if type == "style":
+            style_outputs = outputs[self.NUM_CONTENT_LAYERS :]
+            gram_style_features = [
+                self.gram_matrix(style_output) for style_output in style_outputs
+            ]
+            return gram_style_features
+
+        elif type == "content":
+            content_outputs = outputs[: self.NUM_CONTENT_LAYERS]
+            return content_outputs
+
+        else:
+            style_outputs = outputs[self.NUM_CONTENT_LAYERS :]
+            content_outputs = outputs[: self.NUM_CONTENT_LAYERS]
+            gram_style_features = [
+                self.gram_matrix(style_output) for style_output in style_outputs
+            ]
+            return content_outputs + gram_style_features
+
+    def _loss(self, features, targets, type="style"):
+        """Returns the loss of fearure and target. This is just the mean square error.
+
+        features : list
+            list of features of the content and style images
+        target : list
+            list of features of the content and style images
+        type : str
+            type of loss to calculate. Either "style" or "content"
+        """
+        loss = tf.reduce_mean(tf.square(features - targets))
+        if type == "content":
+            loss = loss * 0.5
+        return loss
+
+    def get_loss(self, features, target, alpha, beta):
+        """Returns the total loss of the style and content images
+
+        Parameters
+        ----------
+        features : list
+            list of features of the content and style images
+        target : list
+            list of features of the content and style images
+        alpha : float
+            weight of the content loss
+        beta : float
+            weight of the style loss
+
+        Returns
+        -------
+        loss : float
+            total loss of the style and content images
+        """
+        style_features = features[self.NUM_CONTENT_LAYERS :]
+        content_features = features[: self.NUM_CONTENT_LAYERS]
+        style_targets = target[self.NUM_CONTENT_LAYERS :]
+        content_targets = target[: self.NUM_CONTENT_LAYERS]
+        style_loss = 0
+        content_loss = 0
+
+        for i in range(self.NUM_STYLE_LAYERS):
+            style_loss += self._loss(style_features[i], style_targets[i], type="style")
+        for i in range(self.NUM_CONTENT_LAYERS):
+            content_loss += self._loss(
+                content_features[i], content_targets[i], type="content"
+            )
+
+        style_loss = beta * style_loss / self.NUM_STYLE_LAYERS
+        content_loss = alpha * content_loss / self.NUM_CONTENT_LAYERS
+        loss = content_loss + style_loss
+        return loss
+
+    def calculate_gradients(self, image, target, alpha, beta):
+        """Calculates the gradients of the loss with respect to the image"""
+        with tf.GradientTape() as tape:
+            features = self.get_features(image, "all")
+            loss = self.get_loss(features, target, alpha, beta)
+        gradients = tape.gradient(loss, image)
+        return gradients, loss
+
+    def update_image(self, image, target, alpha, beta, optimizer):
+        """Updates the image by calculating the gradients and applying them to the image"""
+        gradients, loss = self.calculate_gradients(image, target, alpha, beta)
+        optimizer.apply_gradients([(gradients, image)])
+        image.assign(tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=255.0))
+        return loss
+
+    def load_model(self):
+        """Creates a inception model that returns a list of intermediate output values"""
+        K.clear_session()
+        inception = tf.keras.applications.InceptionV3(
+            include_top=False, weights="imagenet"
+        )
+        inception.trainable = False
+        output_layers = self.content_and_style_layers
+
+        model = tf.keras.models.Model(
+            [inception.input],
+            [inception.get_layer(name).output for name in output_layers],
+        )
+        self.model = model
+        return model
+
+    def stylize_image(
+        self,
+        alpha=1,
+        beta=0.1,
+        epochs=10,
+        steps_per_epoch=10,
+        show_images=True,
+        image_frequency=2,
+        notebook=False,
+        lr=None,
+    ):
+        """Stylizes the image using the style and content images
+
+        Parameters
+        ----------
+        alpha : float, optional
+            Content weight, by default 1
+        beta : float, optional
+            Style weight, by default 0.1
+        epochs : int, optional
+            Number of epochs, by default 10
+        steps_per_epoch : int, optional
+            Number of steps per epoch, by default 10
+        show_images : bool, optional
+            Show images, by default True
+        image_frequency : int, optional
+            Frequency of images to show, by default 2
+        notebook : bool, optional
+            If the code is running on a notebook, by default False
+        lr : float, optional
+            Learning rate, by default None
+
+        Returns
+        -------
+        [PIL.Image]
+            List of images
+        """
+        if self.model is None:
+            K.clear_session()
+            _ = self.load_model()
+        style_image = self.load_image(self.style_image_path)
+        content_image = self.load_image(self.content_image_path)
+
+        style_target = self.get_features(style_image, "style")
+        content_target = self.get_features(content_image, "content")
+
+        target = content_target + style_target
+        image = tf.cast(content_image, dtype=tf.float32)
+        image = tf.Variable(image)
+        # images = []
+        if lr is None:
+            lr = 40.0
+        optimizer = tf.optimizers.Adam(
+            tf.keras.optimizers.schedules.ExponentialDecay(
+                initial_learning_rate=lr, decay_steps=100, decay_rate=0.80
+            )
+        )
+        img = None
+        for epoch in range(epochs):
+            for step in range(steps_per_epoch):
+                loss = self.update_image(image, target, alpha, beta, optimizer)
+                display_image = self.tensor_to_image(image)
+                # images.append(display_image)
+                if show_images:
+                    if (step) % image_frequency == 0:
+                        # save the display_image
+                        display_image.save(f".img.jpg")
+                        # if notebook:
+                        #     display_image = self.tensor_to_image(image)
+                        #     display_fn(
+                        #         display_image,
+                        #         clear=True,
+                        #     )
+                        # else:
+                        #     im = np.array(display_image)
+                        #     if img is None:
+                        #         img = plt.imshow(im)
+                        #     else:
+                        #         img.set_data(im)
+                        #     plt.pause(0.1)
+                        #     plt.draw()
+                        # yield np.array(display_image)
+
+        #         print(f"Epoch: {epoch+1} | Step {step+1} | Loss {loss}", end="\r")
+        # return images