Spaces:
Running
Running
File size: 6,828 Bytes
b16ab70 4381d4f 95110bc 4381d4f b16ab70 eee6b71 95110bc eee6b71 95110bc b16ab70 eee6b71 95110bc eee6b71 b16ab70 95110bc b16ab70 95110bc b16ab70 eee6b71 95110bc b16ab70 eee6b71 95110bc b16ab70 4381d4f b16ab70 4381d4f b16ab70 4381d4f b16ab70 4381d4f b16ab70 95110bc eee6b71 95110bc |
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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
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 *
from models.SRFlow.srflow 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 'srflow_enhanced_image' not in st.session_state:
st.session_state['srflow_enhanced_image'] = None
# Initialize session state for button clicks
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
if 'srflow_clicked' not in st.session_state:
st.session_state['srflow_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":
camera_input = st.camera_input("Take a picture", on_change=lambda: reset_states())
if camera_input is not None:
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['srflow_enhanced_image'] = None
st.session_state['hat_clicked'] = False
st.session_state['rcan_clicked'] = False
st.session_state['srgan_clicked'] = False
st.session_state['srflow_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')
# ------------------------ SRFlow ------------------------ #
if st.button('Enhance with SRFlow'):
with st.spinner('Processing using SRFlow...'):
with st.spinner('Wait for it... the model is processing the image'):
enhanced_image = return_SRFlow_result(image)
st.session_state['srflow_enhanced_image'] = enhanced_image
st.session_state['srflow_clicked'] = True
st.success('Done!')
if st.session_state['srflow_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['srflow_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['srflow_enhanced_image'], 'srflow_enhanced.jpg') |