Initial commit

app.py

@@ -0,0 +1,28 @@
+import os
+
+# Load compressed models from tensorflow_hub
+os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'
+
+import matplotlib as mpl
+
+mpl.rcParams['figure.figsize'] = (12, 12)
+mpl.rcParams['axes.grid'] = False
+
+import tensorflow_hub as hub
+
+from neural_style_transfer import StyleStealer
+
+import gradio as gr
+
+if __name__ == "__main__":
+    hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
+
+    worker = StyleStealer(hub_model)
+
+    demo = gr.Interface(
+        fn=worker.steal,
+        inputs=["image", "image"],
+        outputs=["image"]
+    )
+
+    demo.launch(share=True)

neural_style_transfer.py

@@ -0,0 +1,107 @@
+import os
+
+import tensorflow as tf
+
+os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'
+
+import matplotlib as mpl
+
+mpl.rcParams['figure.figsize'] = (12, 12)
+mpl.rcParams['axes.grid'] = False
+
+import numpy as np
+import PIL.Image
+
+import tensorflow_hub as hub
+
+
+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)
+
+
+def load_img(path_to_img):
+    max_dim = 512
+    img = tf.io.read_file(path_to_img)
+    img = tf.image.decode_image(img, channels=3)
+    img = tf.image.convert_image_dtype(img, tf.float32)
+
+    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
+    long_dim = max(shape)
+    scale = max_dim / long_dim
+
+    new_shape = tf.cast(shape * scale, tf.int32)
+
+    img = tf.image.resize(img, new_shape)
+    img = img[tf.newaxis, :]
+    return img
+
+
+def convert_img(pil_image):
+    max_dim = 512
+
+    # Конвертируем PIL изображение в тензор
+    img = tf.convert_to_tensor(np.array(pil_image))
+
+    # Убедимся, что изображение имеет 3 канала (RGB)
+    if img.shape[-1] != 3:
+        img = tf.stack([img, img, img], axis=-1)
+
+    img = tf.image.convert_image_dtype(img, tf.float32)
+
+    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
+    long_dim = max(shape)
+    scale = max_dim / long_dim
+
+    new_shape = tf.cast(shape * scale, tf.int32)
+
+    img = tf.image.resize(img, new_shape)
+    img = img[tf.newaxis, :]
+    return img
+
+def convert_img_break(pil_image):
+    max_dim = 512
+    # Преобразуем PIL изображение в тензор
+    img = tf.keras.preprocessing.image.img_to_array(pil_image)
+    img = tf.image.convert_image_dtype(img, tf.float32)
+
+    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
+    long_dim = max(shape)
+    scale = max_dim / long_dim
+
+    new_shape = tf.cast(shape * scale, tf.int32)
+
+    img = tf.image.resize(img, new_shape)
+    img = img[tf.newaxis, :]
+    return img
+
+
+class StyleStealer:
+    def __init__(self, model):
+        self.model = model
+
+    def steal_break(self, content_img, style_img):
+        return tensor_to_image(self.model(tf.constant(convert_img_break(content_img)),
+                                          tf.constant(convert_img_break(style_img)))[0])
+
+    def steal(self, content_img, style_img):
+        return tensor_to_image(self.model(tf.constant(convert_img(content_img)),
+                                          tf.constant(convert_img(style_img)))[0])
+
+
+if __name__ == "__main__":
+    content_path = "images/input_4.jpg"
+    style_path = "images/style_6.png"
+
+    content_image = load_img(content_path)
+    style_image = load_img(style_path)
+
+    print("Loading model...")
+    hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
+    print("Model loaded")
+    stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
+    tensor_to_image(stylized_image).save("output.jpg")

requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.28.3
+matplotlib==3.8.4
+numpy==1.26.4
+Pillow==10.3.0
+tensorflow==2.16.1
+tensorflow_hub==0.16.1