Create app.py

app.py

@@ -0,0 +1,171 @@
+import gradio as gr
+import spaces
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import logging
+import time
+from tqdm import tqdm
+
+# Initialize logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("style_transfer_app")
+
+# Set TensorFlow threading options
+tf.config.threading.set_inter_op_parallelism_threads(8)
+tf.config.threading.set_intra_op_parallelism_threads(8)
+
+def load_img(image):
+    """Load and preprocess image for style transfer"""
+    max_dim = 256
+    # Convert PIL Image to tensor
+    img = tf.convert_to_tensor(np.array(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
+
+def vgg_layers(layer_names):
+    """Create VGG model with specified layers"""
+    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
+    vgg.trainable = False
+    outputs = [vgg.get_layer(name).output for name in layer_names]
+    model = tf.keras.Model([vgg.input], outputs)
+    return model
+
+def gram_matrix(input_tensor):
+    """Calculate Gram matrix"""
+    result = 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)
+    return result / num_locations
+
+class StyleContentModel(tf.keras.models.Model):
+    def __init__(self, style_layers, content_layers):
+        super(StyleContentModel, self).__init__()
+        self.vgg = vgg_layers(style_layers + content_layers)
+        self.style_layers = style_layers
+        self.content_layers = content_layers
+        self.num_style_layers = len(style_layers)
+        self.vgg.trainable = False
+
+    def call(self, inputs):
+        inputs = inputs * 255.0
+        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)
+        outputs = self.vgg(preprocessed_input)
+        style_outputs, content_outputs = (outputs[:self.num_style_layers],
+                                        outputs[self.num_style_layers:])
+        style_outputs = [gram_matrix(style_output)
+                        for style_output in style_outputs]
+        content_dict = {content_name: value
+                       for content_name, value
+                       in zip(self.content_layers, content_outputs)}
+        style_dict = {style_name: value
+                     for style_name, value
+                     in zip(self.style_layers, style_outputs)}
+        return {'content': content_dict, 'style': style_dict}
+
+def clip_0_1(image):
+    """Clip tensor values between 0 and 1"""
+    return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)
+
+def style_content_loss(outputs, style_targets, content_targets, style_weight, content_weight):
+    """Calculate style and content loss"""
+    style_outputs = outputs['style']
+    content_outputs = outputs['content']
+    style_loss = tf.add_n([tf.reduce_mean((style_outputs[name]-style_targets[name])**2) 
+                          for name in style_outputs.keys()])
+    style_loss *= style_weight / len(style_outputs)
+    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2) 
+                            for name in content_outputs.keys()])
+    content_loss *= content_weight / len(content_outputs)
+    loss = style_loss + content_loss
+    return loss
+
+@tf.function
+def train_step(image, extractor, style_targets, content_targets, opt, style_weight, content_weight, total_variation_weight):
+    """Perform one training step"""
+    with tf.GradientTape() as tape:
+        outputs = extractor(image)
+        loss = style_content_loss(outputs, style_targets, content_targets, style_weight, content_weight)
+        loss += total_variation_weight * tf.image.total_variation(image)
+    grad = tape.gradient(loss, image)
+    opt.apply_gradients([(grad, image)])
+    image.assign(clip_0_1(image))
+    return loss
+
+def tensor_to_image(tensor):
+    """Convert tensor to PIL Image"""
+    tensor = tensor * 255
+    tensor = np.array(tensor, dtype=np.uint8)
+    if np.ndim(tensor) > 3:
+        tensor = tensor[0]
+    return Image.fromarray(tensor)
+
+# Initialize the style-content model
+style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
+content_layers = ['block5_conv2']
+extractor = StyleContentModel(style_layers, content_layers)
+
+@spaces.GPU(duration=120)  # Style transfer typically needs more than 60s
+def style_transfer_fn(content_image, style_image, progress=gr.Progress(track_tqdm=True)):
+    """Main style transfer function for Gradio interface"""
+    try:
+        # Preprocess images
+        content_img = load_img(content_image)
+        style_img = load_img(style_image)
+
+        # Extract style and content features
+        style_targets = extractor(style_img)['style']
+        content_targets = extractor(content_img)['content']
+        image = tf.Variable(content_img)
+
+        # Set optimization parameters
+        opt = tf.keras.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
+        style_weight = 1e-2
+        content_weight = 1e4
+        total_variation_weight = 30
+
+        epochs = 10
+        steps_per_epoch = 100
+
+        start_time = time.time()
+
+        # Training loop
+        for n in tqdm(range(epochs), desc="Epochs"):
+            for m in tqdm(range(steps_per_epoch), desc="Steps", leave=False):
+                loss = train_step(image, extractor, style_targets, content_targets, 
+                                opt, style_weight, content_weight, total_variation_weight)
+
+        # Convert result to image
+        result_image = tensor_to_image(image)
+        
+        return result_image
+
+    except Exception as e:
+        logger.error(f"Error during style transfer: {e}")
+        raise gr.Error("An error occurred during style transfer.")
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=style_transfer_fn,
+    inputs=[
+        gr.Image(label="Content Image", type="pil"),
+        gr.Image(label="Style Image", type="pil")
+    ],
+    outputs=gr.Image(label="Stylized Image"),
+    title="Neural Style Transfer",
+    description="Upload a content image and a style image to create a stylized combination.",
+    examples=[
+        ["examples/content.jpg", "examples/style.jpg"]
+    ],
+    cache_examples=True
+)
+
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()