Added multiple models

.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

app.py

@@ -1,101 +1,307 @@
 import gradio as gr
 from model import NeuralStyleTransfer
 import tensorflow as tf
+from keras import backend as K
+import numpy as np
 
 
-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,
+def change_dtype_inputs(
+    n_style_layers,
+    n_content_layers,
+    epochs,
+    learning_rate,
+    steps_per_epoch,
+    style_weight,
+    content_weight,
+    var_weight,
 ):
+    return (
+        int(n_style_layers),
+        int(n_content_layers),
+        int(epochs),
+        float(learning_rate),
+        int(steps_per_epoch),
+        float(style_weight),
+        float(content_weight),
+        float(var_weight),
+    )
+
+
+def fit_style_transfer(
+    style_image,
+    content_image,
+    extractor="inception_v3",
+    n_style_layers=2,
+    n_content_layers=3,
+    epochs=4,
+    learning_rate=60.0,
+    steps_per_epoch=100,
+    style_weight=1e-2,
+    content_weight=1e-4,
+    var_weight=1e-12,
+):
+    """
+    Fit the style transfer model to the content and style images.
+
+    Parameters
+    ----------
+
+    style_image: str
+        The path to the style image.
+
+    content_image: str
+        The path to the content image.
+
+    extractor: str
+        The name of the feature extractor to use. Options are
+        "inception_v3", "vgg19", "resnet50", and "mobilenet_v2".
+
+    n_style_layers: int
+        The number of layers to use for the style loss.
+
+    n_content_layers: int
+        The number of layers to use for the content loss.
+
+    epochs: int
+        The number of epochs to train the model for.
+
+    learning_rate: float
+        The learning rate to use for the Adam optimizer.
+
+    steps_per_epoch: int
+        The number of steps to take per epoch.
+
+    style_weight: float
+        The weight to use for the style loss.
+
+    content_weight: float
+        The weight to use for the content loss.
+
+    var_weight: float
+        The weight to use for the total variation loss.
+
+    Returns
+    -------
+    display_image: np.array
+    """
+
+    (
+        n_style_layers,
+        n_content_layers,
+        epochs,
+        learning_rate,
+        steps_per_epoch,
+        style_weight,
+        content_weight,
+        var_weight,
+    ) = change_dtype_inputs(
+        n_style_layers,
+        n_content_layers,
+        epochs,
+        learning_rate,
+        steps_per_epoch,
+        style_weight,
+        content_weight,
+        var_weight,
+    )
+
     model = NeuralStyleTransfer(
-        style_image=style,
-        content_image=content,
+        style_image=style_image,
+        content_image=content_image,
         extractor=extractor,
-        n_content_layers=n_content_layers,
         n_style_layers=n_style_layers,
+        n_content_layers=n_content_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,
+    style_image = model.style_image
+    content_image = model.content_image
+
+    content_and_style_layers = model.get_output_layers()
+
+    # build the model with the layers we need to extract the features from
+    K.clear_session()
+    model.build(content_and_style_layers)
+
+    style_features = model.get_features(style_image, type="style")
+    content_features = model.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
+
+    for epoch in range(epochs):
+        for step in range(steps_per_epoch):
+            losses = model._update_image_with_style(
+                generated_image,
+                style_features,
+                content_features,
+                style_weight,
+                content_weight,
+                optimizer,
+                var_weight,
+            )
+
+            display_image = model.tensor_to_image(generated_image)
+
+            step += 1
+
+            style_loss, content_loss, var_loss = losses
+
+            yield np.array(display_image), style_loss, content_loss, var_loss, epoch, step
+
 
-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.",
-            ),
+
+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))
+
+    extractor = gr.Dropdown(
+        ["inception_v3", "vgg19", "resnet50", "mobilenet_v2"],
+        label="Feature Extractor",
+        value="inception_v3",
+
+    )
+
+    n_content_layers = gr.Slider(
+        1,
+        5,
+        value=3,
+        step=1,
+        label="Content Layers",
+    )
+
+    n_style_layers = gr.Slider(
+        1,
+        5,
+        value=2,
+        step=1,
+        label="Style Layers",
+    )
+
+    epochs = gr.Slider(2, 20, value=4, step=1, label="Epochs")
+
+    learning_rate = gr.Slider(1, 100, value=60, step=1, label="Learning Rate")
+
+    steps_per_epoch = gr.Slider(
+        1,
+        100,
+        value=80,
+        step=1,
+        label="Steps Per Epoch",
+    )
+
+    style_weight = gr.Slider(
+        1e-4,
+        0.5,
+        value=1e-1,
+        step=1e-4,
+        label="Style Weight",
+    )
+
+    content_weight = gr.Slider(
+        1e-3,
+        0.5,
+        value=0.3,
+        step=1e-4,
+        label="Content Weight",
+    )
+
+    var_weight = gr.Slider(
+        0,
+        1e-1,
+        value=1e-5,
+        step=1e-12,
+        label="Total Variation Weight",
+    )
+
+    inputs = [
+        style_image,
+        content_image,
+        extractor,
+        n_style_layers,
+        n_content_layers,
+        epochs,
+        learning_rate,
+        steps_per_epoch,
+        style_weight,
+        content_weight,
+        var_weight,
+    ]
+
+    examples = [
+        [
+            "examples/style_1.jpg",
+            "examples/content_1.jpg",
+            "inception_v3",
+            3,
+            2,
+            4,
+            60,
+            100,
+            1e-2,
+            1e-2,
+            1e-11,
         ],
-        outputs="image",
+        [
+            "examples/style_2.jpg",
+            "examples/content_2.jpg",
+            "inception_v3",
+            3,
+            2,
+            4,
+            60,
+            100,
+            1e-2,
+            1e-2,
+            1e-11,
+        ],
+        [
+            "examples/style_3.jpg",
+            "examples/content_3.jpg",
+            "inception_v3",
+            3,
+            2,
+            4,
+            60,
+            100,
+            1e-2,
+            1e-2,
+            1e-11,
+        ]
+
+    ]
+
+    output_image = gr.Image(type="numpy", label="Output Image", shape=(512, 512))
+
+    style_loss = gr.Number(label="Current Style Loss")
+
+    content_loss = gr.Number(label="Current Content Loss")
+
+    var_loss = gr.Number(label="Current Total Variation Loss")
+
+    curr_epoch = gr.Number(label="Current Epoch")
+
+    curr_step = gr.Number(label="Current Step")
+
+
+
+    outputs = [output_image, style_loss, content_loss, var_loss, curr_epoch, curr_step]
+
+    interface = gr.Interface(
+        fn=fit_style_transfer,
+        inputs=inputs,
+        outputs=outputs,
+        examples=examples
     )
 
-    return demo
+    interface.queue().launch(debug=True)
 
 
-if __name__ == "__main__":
-    demo = hugging_face()
-    demo.launch( )
-    
+main()

model.py

@@ -5,12 +5,24 @@ from keras import backend as K
 
 
 class NeuralStyleTransfer:
-    def __init__(self, style_image, content_image, extractor, n_style_layers=5, n_content_layers=5, display=True):
+    def __init__(self, style_image, content_image, extractor, n_style_layers=5, n_content_layers=5):
         # load the model
         if extractor == "inception_v3":
             self.feature_extractor = tf.keras.applications.InceptionV3(
                 include_top=False, weights="imagenet"
             )
+        elif extractor == "vgg19":
+            self.feature_extractor = tf.keras.applications.VGG19(
+                include_top=False, weights="imagenet"
+            )
+        elif extractor == "resnet50":
+            self.feature_extractor = tf.keras.applications.ResNet50(
+                include_top=False, weights="imagenet"
+            )
+        elif extractor == "mobilenet_v2":
+            self.feature_extractor = tf.keras.applications.MobileNetV2(
+                include_top=False, weights="imagenet"
+            )
         elif isinstance(extractor, tf.keras.Model):
             self.feature_extractor = extractor
         else:
@@ -26,13 +38,6 @@ class NeuralStyleTransfer:
         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)
@@ -215,7 +220,7 @@ class NeuralStyleTransfer:
             ]
         )
         total_loss = style_loss + content_loss
-        return total_loss
+        return total_loss, style_loss, content_loss
 
     def _grad_loss(
         self,
@@ -237,7 +242,7 @@ class NeuralStyleTransfer:
         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(
+            loss, style_loss, content_loss = self._style_content_loss(
                 style_target,
                 style_features,
                 content_target,
@@ -246,9 +251,11 @@ class NeuralStyleTransfer:
                 content_weight,
             )
 
-            loss += var_weight*tf.image.total_variation(generated_image)
+            variational_loss= var_weight*tf.image.total_variation(generated_image)
+
+            loss += variational_loss
         grads = tape.gradient(loss, generated_image)
-        return grads, loss
+        return grads, loss, [style_loss, content_loss, variational_loss]
 
     def _update_image_with_style(
         self,
@@ -260,7 +267,7 @@ class NeuralStyleTransfer:
         optimizer,
         var_weight,
     ):
-        grads, loss = self._grad_loss(
+        grads, loss, loss_list = self._grad_loss(
             generated_image, style_target, content_target, style_weight, content_weight, var_weight
         )
 
@@ -269,84 +276,4 @@ class NeuralStyleTransfer:
         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)
-
-
-    
+        return loss_list