Spaces:
Running
Running
File size: 5,681 Bytes
b16ab70 4381d4f eee6b71 4381d4f b16ab70 eee6b71 b16ab70 eee6b71 b16ab70 eee6b71 b16ab70 eee6b71 b16ab70 4381d4f b16ab70 4381d4f b16ab70 4381d4f b16ab70 4381d4f b16ab70 eee6b71 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import time
import streamlit as st
import numpy as np
from PIL import Image
from io import BytesIO
from models.HAT.hat import *
from models.RCAN.rcan import *
from models.SRGAN.srgan import *
# Initialize session state for enhanced images
if 'hat_enhanced_image' not in st.session_state:
st.session_state['hat_enhanced_image'] = None
if 'rcan_enhanced_image' not in st.session_state:
st.session_state['rcan_enhanced_image'] = None
if 'srgan_enhanced_image' not in st.session_state:
st.session_state['srgan_enhanced_image'] = None
if 'hat_clicked' not in st.session_state:
st.session_state['hat_clicked'] = False
if 'rcan_clicked' not in st.session_state:
st.session_state['rcan_clicked'] = False
if 'srgan_clicked' not in st.session_state:
st.session_state['srgan_clicked'] = False
st.markdown("<h1 style='text-align: center'>Image Super Resolution</h1>", unsafe_allow_html=True)
# Sidebar for navigation
st.sidebar.title("Options")
app_mode = st.sidebar.selectbox("Choose the input source",
["Upload image", "Take a photo"])
# Depending on the choice, show the uploader widget or webcam capture
if app_mode == "Upload image":
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"], on_change=lambda: reset_states())
if uploaded_file is not None:
image = Image.open(uploaded_file).convert("RGB")
elif app_mode == "Take a photo":
# Using JS code to access user's webcam
camera_input = st.camera_input("Take a picture", on_change=lambda: reset_states())
if camera_input is not None:
# Convert the camera image to an RGB image
image = Image.open(camera_input).convert("RGB")
def reset_states():
st.session_state['hat_enhanced_image'] = None
st.session_state['rcan_enhanced_image'] = None
st.session_state['srgan_enhanced_image'] = None
st.session_state['hat_clicked'] = False
st.session_state['rcan_clicked'] = False
st.session_state['srgan_clicked'] = False
def get_image_download_link(img, filename):
"""Generates a link allowing the PIL image to be downloaded"""
# Convert the PIL image to Bytes
buffered = BytesIO()
img.save(buffered, format="PNG")
return st.download_button(
label="Download Image",
data=buffered.getvalue(),
file_name=filename,
mime="image/png"
)
if 'image' in locals():
# st.image(image, caption='Uploaded Image', use_column_width=True)
st.write("")
# ------------------------ HAT ------------------------ #
if st.button('Enhance with HAT'):
with st.spinner('Processing using HAT...'):
with st.spinner('Wait for it... the model is processing the image'):
enhanced_image = HAT_for_deployment(image)
st.session_state['hat_enhanced_image'] = enhanced_image
st.session_state['hat_clicked'] = True
st.success('Done!')
if st.session_state['hat_enhanced_image'] is not None:
col1, col2 = st.columns(2)
col1.header("Original")
col1.image(image, use_column_width=True)
col2.header("Enhanced")
col2.image(st.session_state['hat_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['hat_enhanced_image'], 'hat_enhanced.jpg')
# ------------------------ RCAN ------------------------ #
if st.button('Enhance with RCAN'):
with st.spinner('Processing using RCAN...'):
with st.spinner('Wait for it... the model is processing the image'):
rcan_model = RCAN()
device = torch.device('cpu') if not torch.cuda.is_available() else torch.device('cuda')
rcan_model.load_state_dict(torch.load('models/RCAN/rcan_checkpoint.pth', map_location=device))
enhanced_image = rcan_model.inference(image)
st.session_state['rcan_enhanced_image'] = enhanced_image
st.session_state['rcan_clicked'] = True
st.success('Done!')
if st.session_state['rcan_enhanced_image'] is not None:
col1, col2 = st.columns(2)
col1.header("Original")
col1.image(image, use_column_width=True)
col2.header("Enhanced")
col2.image(st.session_state['rcan_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['rcan_enhanced_image'], 'rcan_enhanced.jpg')
#--------------------------SRGAN--------------------------#
if st.button('Enhance with SRGAN'):
with st.spinner('Processing using SRGAN...'):
with st.spinner('Wait for it... the model is processing the image'):
srgan_model = GeneratorResnet()
device = torch.device('cpu') if not torch.cuda.is_available() else torch.device('cuda')
srgan_model = torch.load('models/SRGAN/srgan_checkpoint.pth', map_location=device)
enhanced_image = srgan_model.inference(image)
st.session_state['srgan_enhanced_image'] = enhanced_image
st.session_state['srgan_clicked'] = True
st.success('Done!')
if st.session_state['srgan_enhanced_image'] is not None:
col1, col2 = st.columns(2)
col1.header("Original")
col1.image(image, use_column_width=True)
col2.header("Enhanced")
col2.image(st.session_state['srgan_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['srgan_enhanced_image'], 'srgan_enhanced.jpg') |