import numpy as np
import tensorflow as tf
import onnxruntime as ort
import gradio as gr
import time
import logging
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_path = "model.onnx"

try:
    start_time = time.perf_counter()
    model = ort.InferenceSession(model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
    logger.info(f"Model successfully loaded from {model_path} in {time.perf_counter()-start_time:.2f} sec")

except Exception as e:
    logger.error(f"Failed to load model from {model_path}: {e}")

def preprocess_image(image):
    "Convert to rgb, normalize and add batch dimension"
    img = Image.open(image).convert("RGB")
    "Warning for image size greater than 512x512"
    MAX_IMAGE_SIZE = 512
    w, h = img.size
    if max(w, h) >= MAX_IMAGE_SIZE:
        logger.warning(f"Large input image {h}x{w}, inference will take longer")

        gr.warning(f"⚠️ Large image detected {h}x{w}, CPU inference will take longer")
    img_arr = np.array(img, dtype=np.float32) / 255.0
    img_arr = np.expand_dims(img_arr, axis=0)

    return img_arr

def super_resolution_image(image):
    try:
        model_inputs = model.get_inputs()[0].name
        model_outputs = model.get_outputs()[0].name
        # Inference
        sr_img = model.run([model_outputs], {model_inputs: image})[0]
        # Postprocess
        sr_img = (np.clip(sr_img, 0, 1)*255).astype(np.uint8)
        return Image.fromarray(sr_img[0])
    
    except Exception as e:
        raise RuntimeError(f"Model inference failed, {str(e)}")
    
def gradio_inference(image):
    img_arr = preprocess_image(image)
    sr_img = super_resolution_image(img_arr)
    return sr_img

demo = gr.Interface(
    fn=gradio_inference,
    inputs=gr.Image(type="filepath"),  
    outputs="image",
    title="EDSR 4x Super-Resolution",
    description="Upscale your images 4× using an onnx edsr model"
)

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