Sa-m's picture
Update app.py
a08c6e3 verified
raw
history blame contribute delete
3.99 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
# Create the Gradio interface with Hugging Face Spaces configuration
demo = gr.Interface(
fn=transform_my_model,
inputs=[
gr.Image(type="numpy", label="Content Image"),
gr.Image(type="numpy", label="Style Image")
],
outputs=gr.Image(type="pil", 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=[
["Content_Images/contnt12.jpg", "VG516.jpg"],
["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
],
cache_examples=True,
allow_flagging=False,
analytics_enabled=False
)
# For Hugging Face Spaces deployment
if __name__ == "__main__":
demo.launch(
share=True # Set to False for Hugging Face Spaces
)