Spaces:
Running
Running
File size: 11,950 Bytes
5bdf4bb |
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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 |
import time
import streamlit as st
import subprocess
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.VDSR.vdsr import *
from models.Interpolation.nearest_neighbor import NearestNeighbor_for_deployment
from models.Interpolation.bilinear import Bilinear_for_deployment
from models.Interpolation.bicubic import Bicubic_for_deployment
subprocess.call('pip install natsort', shell=True)
from models.SRFlow.srflow import *
# Initialize session state for enhanced images
if 'nearest_enhanced_image' not in st.session_state:
st.session_state['nearest_enhanced_image'] = None
if 'bilinear_enhanced_image' not in st.session_state:
st.session_state['bilinear_enhanced_image'] = None
if 'bicubic_enhanced_image' not in st.session_state:
st.session_state['bicubic_enhanced_image'] = None
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
if 'vdsr_enhanced_image' not in st.session_state:
st.session_state['vdsr_enhanced_image'] = None
# Initialize session state for button clicks
if 'nearest_clicked' not in st.session_state:
st.session_state['nearest_clicked'] = False
if 'bilinear_clicked' not in st.session_state:
st.session_state['bilinear_clicked'] = False
if 'bicubic_clicked' not in st.session_state:
st.session_state['bicubic_clicked'] = False
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
if 'vdsr_clicked' not in st.session_state:
st.session_state['vdsr_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['bicubic_enhanced_image'] = None
st.session_state['bilinear_enhanced_image'] = None
st.session_state['nearest_enhanced_image'] = None
st.session_state['vdsr_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
st.session_state['bicubic_clicked'] = False
st.session_state['bilinear_clicked'] = False
st.session_state['nearest_clicked'] = False
st.session_state['vdsr_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("")
# ------------------------ Nearest Neighbor ------------------------ #
if st.button('Enhance with Nearest Neighbor'):
with st.spinner('Processing using Nearest Neighbor...'):
enhanced_image = NearestNeighbor_for_deployment(image)
st.session_state['nearest_enhanced_image'] = enhanced_image
st.session_state['nearest_clicked'] = True
st.success('Done!')
if st.session_state['nearest_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['nearest_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['nearest_enhanced_image'], 'nearest_enhanced.jpg')
# ------------------------ Bilinear ------------------------ #
if st.button('Enhance with Bilinear'):
with st.spinner('Processing using Bilinear...'):
enhanced_image = Bilinear_for_deployment(image)
st.session_state['bilinear_enhanced_image'] = enhanced_image
st.session_state['bilinear_clicked'] = True
st.success('Done!')
if st.session_state['bilinear_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['bilinear_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['bilinear_enhanced_image'], 'bilinear_enhanced.jpg')
# ------------------------ Bicubic ------------------------ #
if st.button('Enhance with Bicubic'):
with st.spinner('Processing using Bicubic...'):
enhanced_image = Bicubic_for_deployment(image)
st.session_state['bicubic_enhanced_image'] = enhanced_image
st.session_state['bicubic_clicked'] = True
st.success('Done!')
if st.session_state['bicubic_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['bicubic_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['bicubic_enhanced_image'], 'bicubic_enhanced.jpg')
# ------------------------ 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')
# ------------------------ VDSR ------------------------ #
if st.button('Enhance with VDSR'):
with st.spinner('Processing using VDSR...'):
# Load the VDSR model
vdsr_model = torch.load('models/VDSR/vdsr_checkpoint.pth', map_location=torch.device('cpu'))
enhanced_image = vdsr_model.inference(image)
st.session_state['vdsr_enhanced_image'] = enhanced_image
st.session_state['vdsr_clicked'] = True
st.success('Done!')
if st.session_state['vdsr_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['vdsr_enhanced_image'], use_column_width=True)
with col2:
get_image_download_link(st.session_state['vdsr_enhanced_image'], 'vdsr_enhanced.jpg') |