SuperResolutionApp / ui /upscaler_ui.py
KuunVo's picture
Add enhanced feature
ca337fc
import streamlit as st
from PIL import Image
import requests
from io import BytesIO
from streamlit_image_comparison import image_comparison
from utils import upscale_image
def ui():
image = None
input_text = None
uploaded_file = None
input_area = st.columns([2, 1, 1])
with input_area[0]:
option = st.selectbox(
"How do you want to provide the image?",
("Fetch from URL", "Upload from local machine")
)
with input_area[2]:
option_scale = st.selectbox(
"Which factors do you want to upscale?",
(2, 3, 4)
)
with input_area[1]:
# , 'SRUNET_interpolation', 'SRUNET_x234_interpolation'
option_model = st.selectbox(
"Which model do you want to use?",
('SRUNET_x2', 'SRUNET_x3', 'SRUNET_x4', 'SRUNET_x234')
)
picture_url_area = st.columns([2, 2], vertical_alignment="top")
with picture_url_area[0]:
if option == "Upload from local machine":
uploaded_file = st.file_uploader(
"Choose an image...", type=["jpg", "jpeg", "png"])
elif option == "Fetch from URL":
input_text = st.text_input("Enter the image URL")
if st.button("Submit"):
if option == "Upload from local machine" and uploaded_file is not None:
try:
image = Image.open(uploaded_file)
# st.image(image, caption="Uploaded Image", use_column_width=True)
except Exception as e:
st.error(f"Error opening image: {e}")
elif option == "Fetch from URL" and input_text:
try:
response = requests.get(input_text)
response.raise_for_status()
image = Image.open(BytesIO(response.content))
# st.image(image, caption="Image from URL", use_column_width=True)
except requests.exceptions.RequestException as e:
st.error(f"Error fetching image: {e}")
if image:
width, height = image.size
if width * int(option_scale) > 1000 or height * int(option_scale) > 1000:
st.error(
"Unable to upscale. The size of upscaled image should be less than 1000x1000")
image = None
if image:
st.header('Results')
# image_resize = image.resize((1024, 1024))
width, height = image.size
picture_url_area[1].text(
f"Image size: {width}x{height} --> {width*int(option_scale)}x{height*int(option_scale)}")
picture_url_area[1].image(
image, caption="Original", use_column_width=True)
img_1 = image.resize(
(width*int(option_scale), height*int(option_scale)), Image.BICUBIC)
img_2 = upscale_image(image, option_model, int(option_scale))
image_comparison(
img1=img_1,
img2=img_2,
)