Create app.py

app.py

@@ -0,0 +1,236 @@
+##############################################################################################################
+# Filename: app.py
+# Description: A Streamlit application to test our implementation of the x4 model,
+# as descirbed in the paper "Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data"
+##############################################################################################################
+#
+# Import libraries.
+#
+import cv2
+import numpy as np
+import requests
+import streamlit as st
+
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from inference.real_esrgan import RealEsrGan
+from io import BytesIO
+from PIL import Image
+
+##############################################################################################################
+
+
+# Function to run inference using the RealEsrGan model.
+def run_inference(
+    uploaded_file,
+    model_name="REALESRGAN_x4",
+    output_path="inferences",
+    upscale=4,
+    extension="auto",
+    device=None,
+    gpu_id=None,
+):
+    try:
+        # Create an RRDBNet model instance.
+        model = RRDBNet(
+            num_in_ch=3,
+            num_out_ch=3,
+            num_feat=64,
+            num_block=23,
+            num_grow_ch=32,
+            scale=upscale,
+        )
+
+        # Set default model path based on the selected model name
+        if model_name == None:
+            model_path = "./models/REALESRGAN_x4.pth"
+        elif model_name == "REALESRGAN_x4":
+            model_path = "./models/REALESRGAN_x4.pth"
+        elif model_name == "REALESRNET_x4":
+            model_path = "./models/REALESRNET_x4.pth"
+
+        # Create an RealEsrGan model instance.
+        upsampler = RealEsrGan(
+            scale=upscale,
+            model_path=model_path,
+            dni_weight=None,
+            model=model,
+            pre_pad=10,
+            half=False,
+            device=device,
+            gpu_id=gpu_id,
+        )
+
+        # Process the input image.
+        if hasattr(
+            uploaded_file, "read"
+        ):  # Check if it's a file uploaded from the local system.
+            img_pil = Image.open(uploaded_file)
+        elif uploaded_file.startswith("http"):  # If it is an image URL.
+            response = requests.get(uploaded_file)
+            img_pil = Image.open(BytesIO(response.content))
+        else:
+            st.warning(
+                "Invalid input. Please provide either an image file or an image URL."
+            )
+            return
+
+        # Convert PIL image to OpenCV format.
+        img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+        # Perform super-resolution using Real-ESRGAN.
+        output, _ = upsampler.enhance(img, upscale=upscale)
+
+        # Determine the file extension for saving the output image.
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = "RGBA"
+            extension = "png"
+        else:
+            img_mode = None
+            if extension == "auto":
+                extension = "png"  # Default extension for images from URL.
+
+        # Save the super resolution image
+        save_path = f"{output_path}/{model_name}_inference.{extension}"
+        cv2.imwrite(save_path, output)
+    except Exception as e:
+        st.error(e)
+    return save_path
+
+
+##############################################################################################################
+
+
+# Function to apply local CSS.
+def local_css(file_name):
+    with open(file_name) as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+
+
+##############################################################################################################
+# Main function to create the Streamlit web application.
+def main():
+    try:
+        # Load CSS.
+        local_css("styles/style.css")
+
+        # Title.
+        title = f"""<p align="center" style="font-family: monospace; color: #FAF9F6; font-size: 2.3rem;">
+                    Super Upscale Resolution with Real-ESRGAN</p>"""
+        st.markdown(title, unsafe_allow_html=True)
+
+        # Toggle button for displaying text input or file uploader.
+        title = f"""<p style="font-family: monospace; color: white;">
+                    Enter Image URL or Upload Image (checkbox):</p>"""
+        st.markdown(title, unsafe_allow_html=True)
+
+        use_image_url = st.checkbox(
+            label="Enter Image URL or Upload Image:", label_visibility="collapsed"
+        )
+
+        # Input for image URL or file uploader based on the checkbox state.
+        if use_image_url:
+            image_url_label = f"""
+                <p style="font-family: monospace; color: white;">Enter Image URL:</p>"""
+            st.markdown(image_url_label, unsafe_allow_html=True)
+
+            image_url = st.text_input(
+                label="Enter Image URL:",
+                value="",
+                label_visibility="collapsed",
+            )
+        else:
+            uploaded_file_label = f"""
+                <p style="font-family: monospace; color: white;">Upload Image:</p>"""
+            st.markdown(uploaded_file_label, unsafe_allow_html=True)
+            uploaded_file = st.file_uploader(
+                label="Upload Image:",
+                type=["jpg", "png", "jpeg"],
+                label_visibility="collapsed",
+            )
+
+        # Dropdown menu for model selection.
+        model_name_label = f"""
+                <p style="font-family: monospace; color: white;">Select Model:</p>"""
+        st.markdown(model_name_label, unsafe_allow_html=True)
+
+        model_name = st.selectbox(
+            label="Select Model:",
+            options=[
+                "REALESRGAN_x4",
+                "REALESRNET_x4",
+            ],
+            label_visibility="collapsed",
+        )
+
+        # Slider for upscale selection.
+        model_name_label = f"""
+                <p style="font-family: monospace; color: white;">Select Upscale Factor. Model works best with x4 upscale:</p>"""
+        st.markdown(model_name_label, unsafe_allow_html=True)
+
+        upscale = st.slider(
+            label="Select Upscale Factor. Model works best with x4 upscale:",
+            min_value=3,
+            max_value=10,
+            value=4,
+            step=1,
+            label_visibility="collapsed",
+        )
+
+        if not use_image_url and uploaded_file is not None:
+            # Image caption.
+            image_caption = f"""<p align="center" style="font-family: monospace; color: #FAF9F6; font-size: 2.3rem;">
+                        Uploaded Image:</p>"""
+            st.markdown(image_caption, unsafe_allow_html=True)
+            st.image(uploaded_file)
+
+        with st.spinner(
+            text="Running Inference. May take up to 3 minutes. Please be patient..."
+        ):
+            if st.button("Run Inference"):
+                if use_image_url and image_url != "":
+                    result_path = run_inference(
+                        uploaded_file=image_url,
+                        model_name=model_name,
+                        upscale=upscale,
+                    )
+                    # Image caption.
+                    image_caption = f"""<p align="center" style="font-family: monospace; color: #FAF9F6; font-size: 2.3rem;">
+                                Resulting Image:</p>"""
+                    st.markdown(image_caption, unsafe_allow_html=True)
+                    st.image(result_path)
+
+                    st.success("Inference completed!")
+                elif not use_image_url and uploaded_file is not None:
+                    result_path = run_inference(
+                        uploaded_file=uploaded_file,
+                        model_name=model_name,
+                        upscale=upscale,
+                    )
+
+                    # Image caption.
+                    image_caption = f"""<p align="center" style="font-family: monospace; color: #FAF9F6; font-size: 2.3rem;">
+                                Resulting Image:</p>"""
+                    st.markdown(image_caption, unsafe_allow_html=True)
+                    st.image(result_path)
+
+                    st.success("Inference completed!")
+                else:
+                    st.warning("Please provide either an image file or an image URL.")
+
+        # GitHub repository of this project.
+        st.markdown(
+            f"""
+                <p align="center" style="font-family: monospace; color: #FAF9F6; font-size: 1rem;">
+                <b>Check out our <a href="https://github.com/GeorgiosIoannouCoder/realesrgan" style="color: #FAF9F6;">GitHub repository</a></b>
+                </p>
+            """,
+            unsafe_allow_html=True,
+        )
+    except Exception as e:
+        st.error(e)
+
+
+##############################################################################################################
+
+if __name__ == "__main__":
+    main()
+##############################################################################################################