Add enhanced feature

app.py

@@ -3,8 +3,9 @@ from ui import upscaler_ui, enhancer_ui
 
 st.set_page_config(layout="wide")
 
+st.header('SUPER RESOLUTION')
 # st.title("Image Upscaler and Enhancer")
-tab1, tab2 = st.tabs(["Upscaler", "Enhancer"])
+tab1, tab2 = st.tabs([  "Upscaler", "Enhancer"]) #
 
 with tab1:
     upscaler_ui.ui()

ui/enhancer_ui.py

@@ -1,2 +1,81 @@
+import streamlit as st
+from PIL import Image
+import requests
+from io import BytesIO
+from streamlit_image_comparison import image_comparison
+from utils import upscale_image, enhanced_image
+
+
 def ui():
-    pass
+    img_selected = None
+    input_text = None
+    uploaded_file = None
+
+    input_image_area = st.columns(2)
+
+    with input_image_area[0]:
+        option = st.selectbox(
+            "How do you want to provide the image?",
+            ("Fetch from URL", "Upload from local machine"),
+            key="option_enhanced"
+        )
+        if option == "Upload from local machine":
+            uploaded_file = st.file_uploader(
+                "Choose an image...", type=["jpg", "jpeg", "png"], key='file_enhanced')
+        elif option == "Fetch from URL":
+            input_text = st.text_input(
+                "Enter the image URL", key='input_enhanced')
+
+        if st.button("Submit", key="submit_enhanced"):
+            if option == "Upload from local machine" and uploaded_file is not None:
+                try:
+                    img_selected = Image.open(uploaded_file)
+                    # st.image(image, caption="Uploaded Image", use_column_width=True)
+                except Exception as e:
+                    st.error(f"Error opening image: {e}")
+            elif option == "Fetch from URL" and input_text:
+                try:
+                    response = requests.get(input_text)
+                    response.raise_for_status()
+                    img_selected = Image.open(BytesIO(response.content))
+                    # st.image(image, caption="Image from URL", use_column_width=True)
+                except requests.exceptions.RequestException as e:
+                    st.error(f"Error fetching image: {e}")
+
+            if img_selected:
+                width, height = img_selected.size
+                if width > 1000 or height > 1000:
+                    st.error(
+                        "Unable to upscale. The size of upscaled image should be less than 1000x1000")
+                    img_selected = None
+
+
+    with input_image_area[1]:
+        option_model = st.selectbox(
+            "Which model do you want to use?",
+            ('SRUNET_x2', 'SRUNET_x3', 'SRUNET_x4', 'SRUNET_x234'),
+            key="option_model_enhanced"
+        )
+
+    
+    if img_selected:
+        st.header('Results')
+        st.text(f'Model: {option_model}')
+        
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(img_selected, caption="Original",
+                        use_column_width=True)
+        with col2:
+            img_enhanced = enhanced_image(img_selected, option_model)
+            # img_enhanced = img_selected.resize((64, 64))
+            col2.image(img_enhanced, caption="Enhanced",
+                    use_column_width=True)
+
+        image_comparison(
+            img1=img_selected,
+            img2=img_enhanced,
+        )
+        
+       
+       

ui/upscaler_ui.py

@@ -62,7 +62,6 @@ def ui():
                     st.error(
                         "Unable to upscale. The size of upscaled image should be less than 1000x1000")
                     image = None
-                    # pass
 
     if image:
         st.header('Results')

utils.py

@@ -73,3 +73,38 @@ def upscale_image(img, model_name, scale_factor):
         img_scale_pred = img_scale_pred.squeeze(0)
     return transforms.ToPILImage()(img_scale_pred).convert(img_mode)
 
+def enhanced_image(img, model_name):
+    img_mode = img.mode
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+
+    transform = transforms.Compose([
+        transforms.ToImage(),
+        transforms.ToDtype(torch.float32, scale=True),
+    ])
+    
+    #Load Model
+    checkpoint = torch.load(get_pretrained_path(
+        model_name), map_location=torch.device('cpu'))
+    model = UNET()
+    model.load_state_dict(checkpoint['best_model_state_dict'])
+    model.eval()
+
+    data = transform(img).clamp(0, 1).unsqueeze(0)
+    h_pad, w_pad = find_padding(data)
+    data = F.pad(data, (0, w_pad, 0, h_pad), mode='reflect')
+    
+
+    with torch.no_grad():
+        img_scale_pred = model(data).clamp(0, 1)
+        if h_pad > 0 and w_pad > 0:
+            img_scale_pred = img_scale_pred[..., :-h_pad, :-w_pad]
+        elif h_pad > 0:
+            img_scale_pred = img_scale_pred[..., :-h_pad, :]
+        elif w_pad > 0:
+            img_scale_pred = img_scale_pred[..., :, :-w_pad]
+        else:
+            img_scale_pred = img_scale_pred
+
+        img_scale_pred = img_scale_pred.squeeze(0)
+    return transforms.ToPILImage()(img_scale_pred).convert(img_mode)