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')