Re-configured SRFlow

app.py

@@ -5,50 +5,54 @@ 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.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"])
+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['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"""
@@ -102,7 +106,8 @@ if 'image' in locals():
         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--------------------------#
+    
+    # --------------------------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'):
@@ -120,4 +125,21 @@ if 'image' in locals():
         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')
+            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')

models/SRFlow/srflow.py

@@ -3,22 +3,23 @@ import torch
 import sys
 sys.path.append('models')
 from SRFlow.code import imread, impad, load_model, t, rgb
+from PIL import Image
+from torchvision.transforms import PILToTensor
 
-
-def return_SRFlow_result(lr_path, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.6):
+def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.6):
     """
     Apply Super-Resolution using SRFlow model to the input LR (low-resolution) image.
 
     Args:
-    - lr_path (str): File path of the LR image. (Refer to OpenCV documentation for suitable image formats with imread)
+    - lr: PIL Image
     - conf_path (str): Configuration file path for the SRFlow model. Default is SRFlow_DF2K_4X.yml.
     - heat (float): Heat parameter for the SRFlow model. Default is 0.6.
 
     Returns:
-    - sr (numpy.ndarray): Super-resolved image in numpy array format.
+    - sr: PIL Image
     """
     model, opt = load_model(conf_path)
-    lr = imread(lr_path)
+    lr = PILToTensor()(lr).permute(1, 2, 0).numpy()
 
     scale = opt['scale']
     pad_factor = 2
@@ -35,5 +36,10 @@ def return_SRFlow_result(lr_path, conf_path='models/SRFlow/code/confs/SRFlow_DF2
     sr = rgb(torch.clamp(sr_t, 0, 1))
     sr = sr[:h * scale, :w * scale]
 
-    sr
+    sr = Image.fromarray((sr).astype('uint8'))
     return sr
+
+if __name__ == '__main__':
+    ip = Image.open('images/demo.png')
+    sr = return_SRFlow_result(ip)
+    print(sr.size)