Sa-m's picture
Update app.py
e1e87af verified
raw
history blame
4.31 kB
# '''NEURAL STYLE TRANSFER '''
# import numpy as np
# import tensorflow as tf
# import tensorflow_hub as hub
# import gradio as gr
# from PIL import Image
# np.set_printoptions(suppress=True)
# model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
# def tensor_to_image(tensor):
# tensor *= 255
# tensor = np.array(tensor, dtype=np.uint8)
# if tensor.ndim > 3:
# tensor = tensor[0]
# return Image.fromarray(tensor)
# def transform_my_model(content_image, style_image):
# content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
# style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
# stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
# return tensor_to_image(stylized_image)
# demo = gr.Interface(
# fn=transform_my_model,
# inputs=[gr.Image(label="Content Image"), gr.Image(label="Style Image")],
# outputs=gr.Image(label="Result"),
# title="Style Transfer",
# examples=[
# ["Content_Images/contnt12.jpg", "VG516.jpg"],
# ["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
# ["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
# ],
# article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin."
# )
# demo.launch(share=True)
'''NEURAL STYLE TRANSFER '''
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import gradio as gr
from PIL import Image
import os
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
np.set_printoptions(suppress=True)
# Load model with error handling
try:
logger.info("Loading TensorFlow Hub model...")
model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
logger.info("Model loaded successfully!")
except Exception as e:
logger.error(f"Error loading model: {str(e)}")
raise
def tensor_to_image(tensor):
try:
tensor *= 255
tensor = np.array(tensor, dtype=np.uint8)
if tensor.ndim > 3:
tensor = tensor[0]
return Image.fromarray(tensor)
except Exception as e:
logger.error(f"Error in tensor_to_image: {str(e)}")
raise
def transform_my_model(content_image, style_image):
try:
if content_image is None or style_image is None:
raise ValueError("Both content and style images are required")
logger.info("Processing images...")
content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
logger.info("Style transfer completed successfully!")
return tensor_to_image(stylized_image)
except Exception as e:
logger.error(f"Error in transform_my_model: {str(e)}")
raise
# Verify example images exist
example_images = [
["Content_Images/contnt12.jpg", "VG516.jpg"],
["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
]
valid_examples = []
for content_path, style_path in example_images:
if os.path.exists(content_path) and os.path.exists(style_path):
valid_examples.append([content_path, style_path])
else:
logger.warning(f"Example image not found: {content_path} or {style_path}")
demo = gr.Interface(
fn=transform_my_model,
inputs=[gr.Image(label="Content Image"), gr.Image(label="Style Image")],
outputs=gr.Image(label="Result"),
title="Neural Style Transfer",
description="Upload a content image and a style image to create a stylized version of your content image.",
examples=valid_examples if valid_examples else None,
article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin."
)
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=True,
show_error=True
)