Update app.py

app.py

@@ -1,107 +1,59 @@
 import gradio as gr
 import numpy as np
-import onnxruntime as ort
-from PIL import Image
 import cv2
+from PIL import Image
 import os
-import requests
-import hashlib
-
-# Configuration - UPDATED MODEL URL
-MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.onnx"
-MODEL_PATH = "realesr.onnx"
-EXPECTED_MD5 = "8a628e89b1e4d9f5f174a3e8c0c7b3b1"  # MD5 hash for verification
-
-def verify_file(file_path, expected_md5):
-    """Verify file integrity using MD5 hash"""
-    if not os.path.exists(file_path):
-        return False
-    with open(file_path, "rb") as f:
-        file_hash = hashlib.md5(f.read()).hexdigest()
-    return file_hash == expected_md5
-
-def download_model():
-    """Download model with verification and retries"""
-    print("Downloading model...")
-    for attempt in range(3):  # Retry up to 3 times
-        try:
-            # Use a mirror URL if primary fails
-            urls = [
-                MODEL_URL,
-                "https://huggingface.co/spaces/akhaliq/Real-ESRGAN/resolve/main/realesr-general-x4v3.onnx"
-            ]
-            
-            for url in urls:
-                try:
-                    response = requests.get(url, stream=True, timeout=30)
-                    response.raise_for_status()
-                    
-                    # Save in chunks
-                    with open(MODEL_PATH, "wb") as f:
-                        for chunk in response.iter_content(chunk_size=8192):
-                            f.write(chunk)
-                    
-                    # Verify download
-                    if verify_file(MODEL_PATH, EXPECTED_MD5):
-                        print("Model downloaded and verified successfully!")
-                        return
-                    os.remove(MODEL_PATH)
-                except Exception:
-                    continue
-            
-            raise ValueError("All download attempts failed")
-        except Exception as e:
-            if attempt == 2:  # Final attempt
-                raise gr.Error(f"Model download failed after 3 attempts: {str(e)}")
-            print(f"Attempt {attempt + 1} failed, retrying...")
 
-# Download model if missing or corrupted
-if not verify_file(MODEL_PATH, EXPECTED_MD5):
-    download_model()
-
-# Initialize ONNX Runtime (CPU only for free tier)
-try:
-    ort_session = ort.InferenceSession(
-        MODEL_PATH,
-        providers=['CPUExecutionProvider']  # Only use CPU on free tier
-    )
-    print("Model loaded successfully!")
-except Exception as e:
-    raise gr.Error(f"Model loading failed: {str(e)}\nTry deleting the space and recreating it.")
-
-def enhance(image):
-    """Image enhancement function"""
+def enhance_image(
+    input_img,
+    contrast=1.2,
+    brightness=10,
+    sharpness=2.0,
+    denoise_strength=10
+):
+    """Enhance image using OpenCV operations"""
     try:
-        # Convert and resize
-        img = np.array(image)
-        if max(img.shape) > 512:  # Free tier memory limit
-            scale = 512 / max(img.shape)
-            img = cv2.resize(img, (0,0), fx=scale, fy=scale)
+        # Convert to OpenCV format
+        img = np.array(input_img)
+        
+        # Contrast and brightness adjustment
+        img = cv2.convertScaleAbs(img, alpha=contrast, beta=brightness)
+        
+        # Denoising
+        img = cv2.fastNlMeansDenoisingColored(img, None, denoise_strength, denoise_strength, 7, 21)
         
-        # Preprocess
-        img = img.astype(np.float32) / 255.0
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        img = np.transpose(img, (2, 0, 1))
-        img = np.expand_dims(img, axis=0)
+        # Sharpening
+        kernel = np.array([[-1,-1,-1], 
+                          [-1,9,-1],
+                          [-1,-1,-1]])
+        img = cv2.filter2D(img, -1, kernel)
         
-        # Inference
-        output = ort_session.run(None, {'input': img})[0]
+        # Color correction
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
+        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_LAB2RGB)
         
-        # Postprocess
-        output = output.squeeze().transpose(1, 2, 0)
-        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-        output = (output * 255).clip(0, 255).astype(np.uint8)
-        return Image.fromarray(output)
+        return Image.fromarray(img)
     except Exception as e:
         raise gr.Error(f"Enhancement failed: {str(e)}")
 
-# Create interface
+# Create interface with adjustable parameters
 demo = gr.Interface(
-    fn=enhance,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Image Enhancement",
-    allow_flagging="never"
+    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"),
+        gr.Slider(0, 20, value=10, label="Denoise Strength")
+    ],
+    outputs=gr.Image(type="pil", label="Enhanced Image"),
+    title="Image Enhancement Tool",
+    allow_flagging="never",
+    examples=["example.jpg"] if os.path.exists("example.jpg") else None
 )
 
 demo.launch()