Upload 3 files

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import tensorflow as tf
+import numpy as np
+import cv2
+import os
+from main import tensor_to_image,StyleContentModel, run_style_transfer
+
+# Define layers and instantiate the model once globally
+content_layers = ['block5_conv2'] 
+style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
+extractor = StyleContentModel(style_layers, content_layers)
+
+def style_transfer_wrapper(content_img_np, style_img_np):
+    """
+    A wrapper to handle I/O for the Gradio interface.
+    Saves numpy arrays to temp files to use with the main function.
+    """
+    if content_img_np is None or style_img_np is None:
+        return None # Return None if either image is missing
+
+    # Save numpy arrays to temporary files
+    content_path = "temp_content.jpg"
+    style_path = "temp_style.jpg"
+    
+    # Gradio provides RGB, but cv2 saves in BGR order
+    cv2.imwrite(content_path, cv2.cvtColor(content_img_np, cv2.COLOR_RGB2BGR))
+    cv2.imwrite(style_path, cv2.cvtColor(style_img_np, cv2.COLOR_RGB2BGR))
+
+    # Run the main process (using fewer iterations for a faster demo)
+    final_tensor = run_style_transfer(content_path, style_path, iterations=500)
+    
+    # Convert tensor to displayable image
+    final_image = tensor_to_image(final_tensor)
+    
+    # Clean up temporary files
+    os.remove(content_path)
+    os.remove(style_path)
+
+    return final_image
+
+## 4. GRADIO UI DEFINITION ##
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎨 Neural Style Transfer")
+    gr.Markdown("Combine the content of one image with the artistic style of another. This demo uses a VGG19 model. Processing can take a minute, especially on CPU.")
+
+    with gr.Row():
+        content_img = gr.Image(label="Content Image", type="numpy", value="https://gradio-builds.s3.amazonaws.com/demo-files/acropolis.jpg")
+        style_img = gr.Image(label="Style Image", type="numpy", value="https://gradio-builds.s3.amazonaws.com/demo-files/starry_night.jpg")
+    
+    run_button = gr.Button("Generate Image", variant="primary")
+    
+    output_img = gr.Image(label="Result")
+    
+    run_button.click(
+        fn=style_transfer_wrapper, 
+        inputs=[content_img, style_img], 
+        outputs=output_img
+    )
+    
+    gr.Markdown("---")
+    gr.Markdown("Based on the paper '[A Neural Algorithm of Artistic Style](https://arxiv.org/abs/1508.06576)' by Gatys et al.")
+
+# Launch the Gradio app
+demo.launch(debug=True)

main.py

@@ -0,0 +1,157 @@
+import tensorflow as tf
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+## 1. UTILITY FUNCTIONS ##
+
+def load_and_process_img(path_to_img, img_size=512):
+    """Loads an image from a path and prepares it for the VGG model."""
+    img = cv2.imread(path_to_img)
+    if img is None:
+        raise FileNotFoundError(f"Could not read image from path: {path_to_img}")
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (img_size, img_size))
+    img = img.astype(np.float32)
+    img = np.expand_dims(img, axis=0)
+    img = tf.keras.applications.vgg19.preprocess_input(img)
+    return tf.convert_to_tensor(img)
+
+def tensor_to_image(tensor):
+    """Converts a tensor back into a displayable image."""
+    tensor = tensor * 255
+    tensor = np.array(tensor, dtype=np.uint8)
+    if np.ndim(tensor) > 3:
+        assert tensor.shape[0] == 1
+        tensor = tensor[0]
+    return tensor
+
+## 2. STYLE TRANSFER MODEL ##
+
+def gram_matrix(input_tensor):
+    """Calculates the Gram matrix of a given tensor."""
+    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):
+    """A custom model to extract style and content features."""
+    def __init__(self, style_layers, content_layers):
+        super(StyleContentModel, self).__init__()
+        # Load the VGG19 model
+        vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
+        vgg.trainable = False
+        
+        # Get the outputs from the specified layers
+        style_outputs = [vgg.get_layer(name).output for name in style_layers]
+        content_outputs = [vgg.get_layer(name).output for name in content_layers]
+        model_outputs = style_outputs + content_outputs
+        
+        # Build the new model
+        self.vgg = tf.keras.Model(vgg.input, model_outputs)
+        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):
+        """Processes an image and returns a dict of style and content features."""
+        # VGG19 expects pixel values in the range of 0-255
+        inputs = inputs * 255.0
+        
+        # Get the outputs from our pre-built VGG model
+        outputs = self.vgg(inputs)
+        
+        # Separate the style and content outputs
+        style_outputs = outputs[:self.num_style_layers]
+        content_outputs = outputs[self.num_style_layers:]
+        
+        # Calculate the Gram matrix for each style layer output
+        style_outputs = [gram_matrix(style_output) for style_output in style_outputs]
+        
+        # Create a dictionary of the features
+        content_dict = {name: value for name, value in zip(self.content_layers, content_outputs)}
+        style_dict = {name: value for name, value in zip(self.style_layers, style_outputs)}
+        
+        return {'content': content_dict, 'style': style_dict}
+
+    
+## 3. MAIN ORCHESTRATION FUNCTION ##
+
+def run_style_transfer(content_path, style_path, iterations=1000, content_weight=1e4, style_weight=1e-2):
+    """
+    Performs the neural style transfer.
+    
+    Returns:
+        A tensor representing the final stylized image.
+    """
+    # Define the layers to use for style and content
+    content_layers = ['block5_conv2'] 
+    style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
+
+    # Build the feature extractor model
+    extractor = StyleContentModel(style_layers, content_layers)
+    
+    # Load content and style images
+    content_image = load_and_process_img(content_path)
+    style_image = load_and_process_img(style_path)
+    
+    # Calculate the target style and content features
+    style_targets = extractor(style_image)['style']
+    content_targets = extractor(content_image)['content']
+    
+    # Initialize the image to be generated and the optimizer
+    generated_image = tf.Variable(content_image)
+    optimizer = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
+
+    @tf.function()
+    def train_step(image):
+        with tf.GradientTape() as tape:
+            outputs = extractor(image)
+            
+            # Calculate total loss
+            content_loss = tf.add_n([tf.reduce_mean((content_targets[name] - outputs['content'][name])**2) for name in content_targets.keys()])
+            style_loss = tf.add_n([tf.reduce_mean((style_targets[name] - outputs['style'][name])**2) for name in style_targets.keys()])
+            total_loss = content_weight * content_loss + style_weight * style_loss
+        
+        # Apply gradients to update the generated image
+        grad = tape.gradient(total_loss, image)
+        optimizer.apply_gradients([(grad, image)])
+        image.assign(tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0))
+
+    # Run the optimization loop
+    for i in range(iterations):
+        train_step(generated_image)
+        if (i + 1) % 100 == 0:
+            print(f"Iteration {i+1}/{iterations}")
+
+    # Post-process to remove VGG preprocessing and return a displayable tensor
+    final_tensor = generated_image
+    final_tensor = tf.reverse(final_tensor, axis=[-1]) # BGR to RGB
+    final_tensor = tf.clip_by_value(final_tensor, 0.0, 255.0)
+    
+    return final_tensor
+
+## 4. SCRIPT EXECUTION ##
+
+if __name__ == '__main__':
+    # Define paths to your images
+    CONTENT_PATH = 'path/to/content_image.jpg'
+    STYLE_PATH = 'path/to/style_image.jpg'
+    
+    # Run the style transfer process
+    final_tensor = run_style_transfer(CONTENT_PATH, STYLE_PATH, iterations=1000)
+    
+    # Convert the final tensor to an image and display it
+    final_image = tensor_to_image(final_tensor)
+    
+    plt.figure(figsize=(8, 8))
+    plt.imshow(final_image)
+    plt.axis('off')
+    plt.title("Stylized Image")
+    plt.show()
+
+    # To save the image
+    final_image_bgr = cv2.cvtColor(final_image, cv2.COLOR_RGB2BGR)
+    cv2.imwrite('stylized_image.png', final_image_bgr)

requirements.txt

@@ -0,0 +1,5 @@
+tensorflow 
+opencv-python-headless
+numpy 
+matplotlib
+gradio