import streamlit as st
from PIL import Image
import numpy as np
import cv2
import torch
from io import BytesIO
from realesrgan import RealESRGANer
from realesrgan.archs.rrdbnet_arch import RRDBNet

# Load Real-ESRGAN model
@st.cache_resource
def load_realesrgan_model():
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
    model_path = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.3.0/RealESRGAN_x4plus.pth'
    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=400, tile_pad=10, pre_pad=0, half=True)
    return upsampler

# Initialize model
upsampler = load_realesrgan_model()

# Title and Description
st.title("Super Clean, Unblur, and Enhance Photo/Image App")
st.markdown("Upload an image to enhance its quality, reduce blurriness, and clean noise.")

# File Upload
uploaded_file = st.file_uploader("Upload an image (JPEG, PNG)", type=["jpeg", "jpg", "png"])

# Process Uploaded Image
if uploaded_file:
    # Load image using PIL
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Convert to NumPy array
    img_array = np.array(image)

    # Enhancement Options
    operation = st.selectbox("Choose an Operation", ["Super Clean", "Unblur and Enhance"])

    if st.button("Process Image"):
        # Perform Super Cleaning
        if operation == "Super Clean":
            try:
                result, _ = upsampler.enhance(img_array, outscale=4)
                output_image = Image.fromarray(result)
                st.image(output_image, caption="Super Cleaned Image", use_container_width=True)
            except Exception as e:
                st.error(f"Error during super cleaning: {e}")

        # Perform Unblur and Enhance
        elif operation == "Unblur and Enhance":
            # Apply unsharp mask
            gaussian = cv2.GaussianBlur(img_array, (9, 9), 10.0)
            unblurred = cv2.addWeighted(img_array, 1.5, gaussian, -0.5, 0)

            # Enhance colors using CLAHE (Contrast Limited Adaptive Histogram Equalization)
            lab = cv2.cvtColor(unblurred, cv2.COLOR_RGB2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
            cl = clahe.apply(l)
            enhanced = cv2.merge((cl, a, b))
            final_img = cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB)

            output_image = Image.fromarray(final_img)
            st.image(output_image, caption="Unblurred and Enhanced Image", use_container_width=True)

        # Provide Download Option
        buf = BytesIO()
        output_image.save(buf, format="PNG")
        byte_im = buf.getvalue()
        st.download_button(label="Download Processed Image", data=byte_im, file_name="processed_image.png", mime="image/png")