Configure for RCAN

app.py

@@ -4,6 +4,8 @@ import numpy as np
 from PIL import Image
 from io import BytesIO
 from models.HAT.hat import *
+from models.RCAN.rcan import *
+
 # Initialize session state for enhanced images
 if 'hat_enhanced_image' not in st.session_state:
     st.session_state['hat_enhanced_image'] = None
@@ -55,49 +57,39 @@ 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'):
-                # Simulate a delay for processing image
-                                
+        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!')
-        # Display the low and high resolution images side by side  
     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'):
-                # Simulate a delay for processing image
-                time.sleep(2)  # replace this with actual model processing code
-                
-                enhanced_image = image 
-                # Display the low and high resolution images side by side
+                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')
-    
-
-
-