templates/index.html

import streamlit as st
from PIL import Image
import torch
from RealESRGAN import RealESRGAN
from io import BytesIO
import base64
import streamlit.components.v1 as components

# Function to load the model based on scale and anime toggle
def load_model(scale, anime=False):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = RealESRGAN(device, scale=scale, anime=anime)
    model_path = {
        (2, False): 'model/RealESRGAN_x2.pth',
        (4, False): 'model/RealESRGAN_x4plus.pth',
        (8, False): 'model/RealESRGAN_x8.pth',
        (4, True): 'model/RealESRGAN_x4plus_anime_6B.pth'
    }[(scale, anime)]
    model.load_weights(model_path)
    return model

def enhance_image(image, scale, anime):
    model = load_model(scale, anime=anime)
    sr_image = model.predict(image)
    
    buffer = BytesIO()
    sr_image.save(buffer, format="PNG")
    buffer.seek(0)
    return sr_image, buffer

def get_base64_image(image):
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def image_comparison_slider(original_base64, enhanced_base64):
    slider_html = f"""
    <style>

    .img-comp-container {{

        position: relative;

        width: 100%;

        max-width: 800px;

        margin: auto;

    }}

    .img-comp-img {{

        position: absolute;

        width: 100%;

        height: 100%;

        overflow: hidden;

    }}

    .img-comp-img img {{

        width: 100%;

        height: auto;

        display: block;

    }}

    .img-comp-overlay {{

        position: absolute;

        top: 0;

        left: 0;

        width: 50%; /* Start at 50% for better default view */

        height: 100%;

        overflow: hidden;

        background-color: rgba(0,0,0,0.5);

        transition: 0.4s;

    }}

    .img-comp-slider {{

        position: absolute;

        cursor: ew-resize;

        width: 40px;

        height: 40px;

        background-color: #2196F3;

        border-radius: 50%;

        transform: translateY(-50%);

        top: 50%;

        left: 50%;

        margin-left: -20px;

        margin-top: -20px;

        z-index: 9;

        opacity: 0.7;

        transition: 0.3s;

    }}

    .img-comp-slider:hover {{

        opacity: 1;

    }}

    </style>
    <div class="img-comp-container">
        <div class="img-comp-img">
            <img src="data:image/png;base64,{original_base64}" />
        </div>
        <div class="img-comp-img">
            <img src="data:image/png;base64,{enhanced_base64}" />
        </div>
        <div class="img-comp-overlay"></div>
        <div class="img-comp-slider"></div>
    </div>
    <script>

    function initComparisons() {{

        var x, i;

        x = document.getElementsByClassName("img-comp-container");

        for (i = 0; i < x.length; i++) {{

            compareImages(x[i]);

        }}

        function compareImages(img) {{

            var slider, overlay, clicked = 0, w, h;

            w = img.offsetWidth;

            h = img.offsetHeight;

            slider = img.getElementsByClassName("img-comp-slider")[0];

            overlay = img.getElementsByClassName("img-comp-overlay")[0];

            img.addEventListener("mousemove", slide);

            img.addEventListener("touchmove", slide);

            img.addEventListener("mousedown", setActive);

            window.addEventListener("mouseup", removeActive);

            function slide(e) {{

                if (clicked === 0) return false;

                e.preventDefault();

                var pos = getCursorPos(e);

                if (pos.x > w) pos.x = w;

                if (pos.x < 0) pos.x = 0;

                slider.style.left = pos.x + "px";

                overlay.style.width = pos.x + "px";

            }}

            function getCursorPos(e) {{

                var a, x = 0;

                e = (e.changedTouches) ? e.changedTouches[0] : e;

                a = img.getBoundingClientRect();

                x = e.pageX - a.left;

                x = x - window.pageXOffset;

                return {{ x: x }};

            }}

            function setActive(e) {{

                clicked = 1;

                e.preventDefault();

            }}

            function removeActive() {{

                clicked = 0;

            }}

        }}

    }}

    initComparisons();

    </script>
    """
    components.html(slider_html, height=600, scrolling=True)



def main():
    st.title("Generative AI Image Restoration")

    # Image upload
    uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
    
    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption="Original Image", use_column_width=True)
        
        # Anime toggle
        anime = st.checkbox("Anime Image", value=False)
        
        # Conditional scale options
        if anime:
            scale = "4x"  # Set to 4x automatically when anime is selected
        else:
            scale = st.radio("Upscaling Factor", ["2x", "4x", "8x"], index=0)
        
        scale_value = int(scale.replace('x', ''))
        
        # Enhance button
        if st.button("Restore Image"):
            enhanced_image, buffer = enhance_image(image, scale_value, anime)
            
            # Convert images to base64 for comparison slider
            original_base64 = get_base64_image(image)
            enhanced_base64 = get_base64_image(enhanced_image)
            
            # Show comparison slider
            image_comparison_slider(original_base64, enhanced_base64)
            
            # Download button
            st.download_button(
                label="Download Enhanced Image",
                data=buffer,
                file_name="enhanced_image.png",
                mime="image/png"
            )

if __name__ == "__main__":
    main()