Update app.py

app.py

@@ -1,37 +1,49 @@
 import gradio as gr
-import cv2
-import numpy as np
+from super_image import EdsrModel, ImageLoader
 from PIL import Image
+import requests
+import io
 
-def enhance_image(input_img):
-    try:
-        # Convert to numpy array if needed
-        if isinstance(input_img, Image.Image):
-            input_img = np.array(input_img)
-        
-        # Initialize OpenCV super-resolution
-        sr = cv2.dnn_superres.DnnSuperResImpl_create()
-        
-        # Load a lightweight model (e.g., EDSR x2)
-        # Note: You'll need to add "EDSR_x2.pb" to your Space!
-        sr.readModel("EDSR_x2.pb")
-        sr.setModel("edsr", 2)  # 2x scaling
-        
-        # Upscale
-        enhanced_img = sr.upsample(input_img)
-        
-        return Image.fromarray(enhanced_img)
-    except Exception as e:
-        print(f"Error: {e}")
-        return input_img  # Fallback to original
+def enhance_image(input_image, scale: int = 2):
+    # Convert Gradio input to PIL Image
+    if isinstance(input_image, str):  # If it's a file path
+        image = Image.open(input_image)
+    else:  # If it's a numpy array from Gradio
+        image = Image.fromarray(input_image)
+
+    # Load the pre-trained EDSR model
+    model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=scale)
+    
+    # Process the image
+    inputs = ImageLoader.load_image(image)
+    preds = model(inputs)
+
+    # Return the enhanced image
+    return ImageLoader.save_image(preds, output_path=None)  # Returns PIL Image
 
 # Gradio Interface
-iface = gr.Interface(
-    fn=enhance_image,
-    inputs=gr.Image(label="Upload Image"),
-    outputs=gr.Image(label="Enhanced"),
-    title="🚀 CPU Image Enhancer",
-    description="Upscales images 2x using OpenCV (runs on CPU)."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## 🚀 EDSR Image Upscaler (2x, 3x, 4x)")
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload Image", type="pil")
+            scale = gr.Dropdown(choices=[2, 3, 4], value=2, label="Upscale Factor")
+            submit_btn = gr.Button("Enhance")
+        with gr.Column():
+            output_image = gr.Image(label="Enhanced Image")
+    
+    submit_btn.click(
+        fn=enhance_image,
+        inputs=[input_image, scale],
+        outputs=output_image
+    )
+
+    gr.Markdown("""
+    **How it works:**  
+    - Uses EDSR (Enhanced Deep Super-Resolution) base model  
+    - Supports 2x, 3x, and 4x upscaling  
+    - Runs on CPU (may be slow for large images)  
+    - Try with images <1000px for best performance  
+    """)
 
-iface.launch()
+demo.launch()