File size: 1,209 Bytes
883346a |
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 |
import streamlit as st
from PIL import Image
import torch
from model import ModelColorization
from utils import process_gs_image, inverse_transform_cs
# create the model
model = ModelColorization().from_pretrained("sebastiansarasti/AutoEncoderImageColorization")
# create the streamlit app
st.title("Image Colorization App")
st.write("This is an app to colorize black and white images.")
# create a botton to upload the image
uploaded_file = st.file_uploader("Choose an image...", type="jpg")
# check if the image is uploaded
if uploaded_file is not None:
# display the image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image.", use_container_width=True)
# create a button to colorize the image
if st.button("Colorize"):
# process the grayscale image
image, original_size = process_gs_image(image)
# run the model
model.eval()
with torch.no_grad():
result = model(image)
# colorize the image
colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
# display the colorized image
st.image(colorized_image, caption="Colorized Image.", use_container_width=True) |