app.py

import gradio as gr
from model import NeuralStyleTransfer
import tensorflow as tf
from keras import backend as K
import numpy as np


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,
    n_style_layers,
    n_content_layers,
    epochs,
    learning_rate,
    steps_per_epoch,
    style_weight,
    content_weight,
    var_weight,
):
    """
    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_image,
        content_image=content_image,
        extractor=extractor,
        n_style_layers=n_style_layers,
        n_content_layers=n_content_layers,
    )

    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.50
        )
    )

    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 main():
    content_image = gr.Image(type="filepath", label="Content Image")
    style_image = gr.Image(type="filepath", label="Style Image")

    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=3,
        step=1,
        label="Style Layers",
    )

    epochs = gr.Slider(2, 20, value=4, step=1, label="Epochs")

    learning_rate = gr.Slider(1, 50, value=20, step=1, label="Learning Rate")

    steps_per_epoch = gr.Slider(
        1,
        20,
        value=5,
        step=1,
        label="Steps Per Epoch",
    )

    style_weight = gr.Slider(
        0,
        1,
        value=0.1,
        label="Style Weight",
    )

    content_weight = gr.Slider(
        0,
        1,
        value=1,
        label="Content Weight",
    )

    var_weight = gr.Slider(
        0,
        0.8,
        value=0,
        label="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/van_gogh.jpg",
            "examples/scene.jpg",
            "inception_v3",
            5,
            5,
            10,
            40,
            10,
            0.1,
            1,
            0.0,
        ],
        [
            "examples/painting.jpg",
            "examples/swan.jpg",
            "vgg19",
            5,
            5,
            10,
            40,
            10,
            0.1,
            1,
            0.0,
        ]
    ]

    output_image = gr.Image(type="numpy", label="Output Image")

    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")

    title = "Neural Style Transfer"

    description = """### This app uses a neural network to transfer the style of one image to another. \n### The `style image` is the image whose style you want to transfer, and the `content image` is the image you want to transfer the style to. \n### The `feature extractor` is the neural network used to extract the features from the images. \n### The number of `style layers` and `content layers` are the number of layers in the feature extractor used to extract the style and content features respectively. \n### The `epochs`, `learning_rate`, `steps_per_epoch`, `style_weight`, `content_weight`, and `total_variation_weight` are all **hyperparameters** that affect the style transfer process. \n### The style weight controls how much the style image affects the output image, the content weight controls how much the content image affects the output image, and the variation weight controls how much the total variation of the output image affects the output image. The total variation of an image is the sum of the absolute differences for neighboring pixel-values in the image. The total variation loss is used to smooth the output image. The higher the variation weight, the smoother the output image will be."""



    outputs = [output_image, style_loss, content_loss, var_loss, curr_epoch, curr_step]

    interface = gr.Interface(
        fn=fit_style_transfer,
        inputs=inputs,
        outputs=outputs,
        title=title,
        description=description,
        examples=examples,
        theme='gradio/monochrome' 
        
    )

    interface.queue().launch(server_name="0.0.0.0", server_port=7860)

main()