import gradio as gr
import numpy as np
import tensorflow as tf
import keras
from huggingface_hub import from_pretrained_keras
from PIL import Image
import io
import gc

# Load MIRNet model from Hugging Face
model = from_pretrained_keras("keras-io/lowlight-enhance-mirnet", compile=False)

# TensorFlow graph mode for performance
@tf.function
def enhance_image(img, passes):
    for _ in tf.range(passes):
        img = model(img)
    return img

# Inference function
def process_image(input_img: Image.Image, passes: int):
    try:
        # Convert to RGB, Resize, Normalize
        input_img = input_img.convert("RGB").resize((256, 256), Image.LANCZOS)
        img_array = keras.preprocessing.image.img_to_array(input_img).astype("float32") / 255.0
        img_array = np.expand_dims(img_array, axis=0)

        # Enhance
        output = enhance_image(tf.convert_to_tensor(img_array), passes)
        enhanced_img = (output[0].numpy() * 255.0).clip(0, 255).astype('uint8')
        result_img = Image.fromarray(enhanced_img, "RGB")

        # Return both images for preview
        return result_img
    finally:
        # Memory cleanup
        del img_array, output, enhanced_img
        gc.collect()

# Gradio Interface
title = "🌃 Low-Light Image Enhancer"
description = """
Boost visibility of dark images using deep learning (MIRNet)<br>
Built for Bharatiya Antariksh Hackathon 2025 🚀 – Team <strong>CodeKarma</strong>
"""

demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="📷 Upload Low-Light Image (JPG/PNG)"),
        gr.Slider(minimum=1, maximum=3, value=1, step=1, label="🔁 Enhancement Passes")
    ],
    outputs=[
        gr.Image(type="pil", label="✨ Enhanced Image")
    ],
    title=title,
    description=description,
    allow_flagging="never",
    theme="soft",
)

if __name__ == "__main__":
    demo.launch()