Interpolation added (#5)

- Interpolation added (ff4d653f5d182fc8825c2bf92bb9e4a65905c97f)


Co-authored-by: Nguyễn Nam <namnguyen2103@users.noreply.huggingface.co>

app.py

@@ -7,11 +7,20 @@ from io import BytesIO
 from models.HAT.hat import *
 from models.RCAN.rcan import *
 from models.SRGAN.srgan 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:
@@ -22,6 +31,12 @@ if 'srflow_enhanced_image' not in st.session_state:
     st.session_state['srflow_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:
@@ -52,10 +67,16 @@ def reset_states():
     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['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
     
 def get_image_download_link(img, filename):
     """Generates a link allowing the PIL image to be downloaded"""
@@ -72,6 +93,53 @@ def get_image_download_link(img, filename):
 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'):

models/Interpolation/bicubic.py

@@ -0,0 +1,7 @@
+from PIL import Image
+from torchvision import transforms
+
+def Bicubic_for_deployment(lr_image):
+    w, h = lr_image.size
+    sr_image = transforms.functional.resize(lr_image, size=(h*4, w*4), interpolation=transforms.InterpolationMode.BICUBIC, antialias=True)
+    return sr_image

models/Interpolation/bilinear.py

@@ -0,0 +1,7 @@
+from PIL import Image
+from torchvision import transforms
+
+def Bilinear_for_deployment(lr_image):
+    w, h = lr_image.size
+    sr_image = transforms.functional.resize(lr_image, size=(h*4, w*4), interpolation=transforms.InterpolationMode.BILINEAR, antialias=True)
+    return sr_image

models/Interpolation/nearest_neighbor.py

@@ -0,0 +1,7 @@
+from PIL import Image
+from torchvision import transforms
+
+def NearestNeighbor_for_deployment(lr_image):
+    w, h = lr_image.size
+    sr_image = transforms.functional.resize(lr_image, size=(h*4, w*4),interpolation=transforms.InterpolationMode.NEAREST,antialias=False)
+    return sr_image