File size: 2,737 Bytes
8b06175 9bf033b 8b06175 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import streamlit as st
import os
from PIL import Image
from torchvision.models import vgg19
from model import StyleTransferModel
from trainer import trainer_fn
from utils import process_image, tensor_to_image
base_model = vgg19(pretrained=True).features
final_model = StyleTransferModel(base_model)
# define the title of the app
st.title('Style Transfer App')
# define the description of the app
st.write('This app applies the style of one image to another image. This can be used to create artistic images.')
# get all image files in the 'styles' folder
image_files = [f for f in os.listdir('styles') if f.lower().endswith(('png', 'jpg', 'jpeg', 'gif', 'bmp'))]
# display the images
st.write('Select style art to apply into your image:')
# Check how many images are available and set columns accordingly
num_images = len(image_files)
cols = st.columns(num_images)
# Define the size to which the images will be resized (width, height)
resize_width = 300
resize_height = 300
# show each image in a corresponding column
for idx, img_file in enumerate(image_files):
with cols[idx]:
st.write(f"Style {idx + 1}")
img_path = f'styles/{img_file}'
img = Image.open(img_path)
# Redimensionar la imagen
img_resized = img.resize((resize_width, resize_height))
st.image(img_resized, use_container_width=True)
# create a file uploader for the content image
st.write('Upload the content image:')
content_image = st.file_uploader('Content Image', type=['jpg', 'jpeg'])
# create the botton to select the style image between 1, 2, and 3
choice = st.selectbox('Select the style art:', [f'Style {i + 1}' for i in range(num_images)])
# create a button to run the model
if st.button('Apply Style Transfer'):
if content_image is not None:
# get the content image
content_img = Image.open(content_image)
# get the style image
style_choice = choice.split()[-1] # Extract style number from "Style 1", "Style 2", etc.
style_img = Image.open(os.path.join('styles', image_files[int(style_choice) - 1])) # Get full path
# preprocess the images
content_img = process_image(content_img)
style_img = process_image(style_img)
# run the model
st.write('Applying Style Transfer...')
target_image = trainer_fn(
content_img, style_img, content_img.clone().requires_grad_(True), final_model
)
# convert the tensor to image
target_image = tensor_to_image(target_image.squeeze(0))
# display the result
st.write('Result:')
st.image(target_image, use_container_width=True)
else:
st.write('Please upload a content image')
|