init space

.gitattributes

@@ -33,3 +33,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,3 @@
+.idea/
+__pycache__/
+get.py

app.py

@@ -0,0 +1,57 @@
+import gradio as gr
+from style_transfer import StyleTransfer
+
+style = StyleTransfer()
+
+
+def predict(content_image, style_image):
+    return style.transfer(content_image, style_image)
+
+
+footer = r"""
+<center>
+<b>
+Demo for <a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization'>Style Transfer</a>
+</b>
+</center>
+"""
+
+coffe = r"""
+<center>
+<a href="https://www.buymeacoffee.com/leonelhs"> <img 
+src="https://img.buymeacoffee.com/button-api/?text=Buy me a 
+coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000
+&coffee_colour=ffffff" /></a>
+</center>
+"""
+
+with gr.Blocks(title="Style Transfer") as app:
+    gr.HTML("<center><h1>Style Transfer</h1></center>")
+    gr.HTML("<center><h3>Fast Style Transfer for Arbitrary Styles</h3></center>")
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            content_img = gr.Image(type="filepath", label="Content image")
+            style_img = gr.Image(type="filepath", label="Style image")
+            run_btn = gr.Button(variant="primary")
+        with gr.Column():
+            output_img = gr.Image(type="pil", label="Output image")
+            gr.ClearButton(components=[content_img, style_img, output_img], variant="stop")
+
+    run_btn.click(predict, [content_img, style_img], [output_img])
+
+    with gr.Row():
+        blobs_c = [[f"examples/contents/{x:02d}.jpg"] for x in range(1, 4)]
+        examples_c = gr.Dataset(components=[content_img], samples=blobs_c)
+        examples_c.click(lambda x: x[0], [examples_c], [content_img])
+    with gr.Row():
+        blobs_s = [[f"examples/styles/{x:02d}.jpg"] for x in range(1, 12)]
+        examples_s = gr.Dataset(components=[style_img], samples=blobs_s)
+        examples_s.click(lambda x: x[0], [examples_s], [style_img])
+
+    with gr.Row():
+        gr.HTML(footer)
+    with gr.Row():
+        gr.HTML(coffe)
+
+app.launch(share=False, debug=True, show_error=True)
+app.queue()

style_transfer.py

@@ -0,0 +1,63 @@
+#############################################################################
+#
+#   Source from:
+#   https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization
+#   Forked from:
+#   Reimplemented by: Leonel Hernández
+#
+##############################################################################
+
+import PIL.Image
+import numpy as np
+import tensorflow as tf
+from huggingface_hub import snapshot_download
+
+print("TF Version: ", tf.__version__)
+print("Eager mode enabled: ", tf.executing_eagerly())
+print("GPU available: ", tf.config.list_physical_devices('GPU'))
+
+STYLE_REPO_ID = "leonelhs/arbitrary-image-stylization-v1"
+
+
+def crop_center(image):
+    """Returns a cropped square image."""
+    shape = image.shape
+    new_shape = min(shape[1], shape[2])
+    offset_y = max(shape[1] - shape[2], 0) // 2
+    offset_x = max(shape[2] - shape[1], 0) // 2
+    image = tf.image.crop_to_bounding_box(
+        image, offset_y, offset_x, new_shape, new_shape)
+    return image
+
+
+def load_image(image_path, image_size=(256, 256)):
+    """Loads and preprocesses images."""
+    img = tf.io.decode_image(
+        tf.io.read_file(image_path),
+        channels=3, dtype=tf.float32)[tf.newaxis, ...]
+    img = crop_center(img)
+    img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)
+    return img
+
+
+def tensor_to_image(tensor):
+    tensor = tensor * 255
+    tensor = np.array(tensor, dtype=np.uint8)
+    if np.ndim(tensor) > 3:
+        assert tensor.shape[0] == 1
+        tensor = tensor[0]
+    return PIL.Image.fromarray(tensor)
+
+
+class StyleTransfer:
+
+    def __init__(self):
+        model_path = snapshot_download(STYLE_REPO_ID)
+        self.model = tf.saved_model.load(model_path)
+
+    def transfer(self, content_image, style_image):
+        content_image = load_image(content_image, (384, 384))
+        style_image = load_image(style_image, (256, 256))
+        style_image = tf.nn.avg_pool(style_image, ksize=[3, 3], strides=[1, 1], padding='SAME')
+        stylized_image = self.model(tf.constant(content_image), tf.constant(style_image))
+        return tensor_to_image(stylized_image[0])