Update app.py

app.py

@@ -1,45 +1,48 @@
 import gradio as gr
 import numpy as np
-import cv2
-import torch
+import onnxruntime as ort
 from PIL import Image
+import cv2
+import os
 
-# Simple image enhancement function (no complex dependencies)
-def enhance_image(input_img, contrast=1.2, brightness=10, sharpness=2.0):
-    # Convert to OpenCV format
-    img = np.array(input_img)
-    
-    # Contrast and brightness adjustment
-    img = cv2.convertScaleAbs(img, alpha=contrast, beta=brightness)
-    
-    # Sharpening
-    kernel = np.array([[-1,-1,-1], 
-                       [-1,9,-1],
-                       [-1,-1,-1]])
-    img = cv2.filter2D(img, -1, kernel)
-    
-    # Color correction
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
-    l, a, b = cv2.split(img)
-    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
-    l = clahe.apply(l)
-    img = cv2.merge((l,a,b))
-    img = cv2.cvtColor(img, cv2.COLOR_LAB2BGR)
+# Load ONNX model
+model_path = "esrgan.onnx"  # Replace with your ONNX file name
+ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+def preprocess(img):
+    """Convert PIL image to ONNX-compatible input"""
+    img = np.array(img)
+    img = img.astype(np.float32) / 255.0  # Normalize to [0,1]
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # ESRGAN expects BGR
+    img = np.transpose(img, (2, 0, 1))  # HWC to CHW
+    img = np.expand_dims(img, axis=0)  # Add batch dimension
+    return img
+
+def postprocess(output):
+    """Convert model output to PIL image"""
+    output = output.squeeze()  # Remove batch dim
+    output = np.transpose(output, (1, 2, 0))  # CHW to HWC
+    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)  # BGR to RGB
+    output = (output * 255.0).clip(0, 255).astype(np.uint8)  # Denormalize
+    return Image.fromarray(output)
+
+def enhance(image):
+    # Resize if too large (free-tier GPU memory is limited)
+    if max(image.size) > 1024:
+        image = image.resize((512, 512))
     
-    return Image.fromarray(img)
+    # Preprocess → Inference → Postprocess
+    input_tensor = preprocess(image)
+    output = ort_session.run(None, {'input': input_tensor})[0]
+    return postprocess(output)
 
-# Gradio interface with adjustable parameters
+# Gradio Interface
 demo = gr.Interface(
-    fn=enhance_image,
-    inputs=[
-        gr.Image(type="pil", label="Input Image"),
-        gr.Slider(0.5, 2.0, value=1.2, label="Contrast"),
-        gr.Slider(0, 30, value=10, label="Brightness"),
-        gr.Slider(0.5, 3.0, value=2.0, label="Sharpness")
-    ],
-    outputs=gr.Image(type="pil", label="Enhanced Image"),
-    title="Professional Image Enhancement",
-    examples=["example.jpg"] if os.path.exists("example.jpg") else None
+    fn=enhance,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.Image(type="pil", label="Enhanced"),
+    title="ESRGAN Image Enhancement (ONNX)",
+    examples=["example.jpg"] if os.path.exists("example.jpg") else None,
 )
 
 demo.launch()