Update app.py

app.py

@@ -4,6 +4,8 @@ import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
 import numpy as np
+import cv2
+import tempfile
 
 # ---------------------------
 # MODEL ARCHITECTURE
@@ -60,7 +62,7 @@ model.eval()
 # ---------------------------
 
 transform = transforms.Compose([
-    transforms.Resize((128,128)),
+    transforms.Resize((128, 128)),
     transforms.ToTensor()
 ])
 
@@ -68,13 +70,9 @@ transform = transforms.Compose([
 # INFERENCE FUNCTION
 # ---------------------------
 
-import gradio as gr
-import numpy as np
-from PIL import Image
-import tempfile
-
 def enhance_image(input_image):
     img = input_image.convert("RGB")
+    original_size = img.size
 
     input_tensor = transform(img).unsqueeze(0).to(device)
 
@@ -84,19 +82,37 @@ def enhance_image(input_image):
     output_img = output.squeeze().permute(1,2,0).cpu().numpy()
     output_img = (output_img * 255).astype(np.uint8)
 
-    # 🔥 SAVE TEMP FILE FOR DOWNLOAD
+    # Resize back to original size
+    output_img = Image.fromarray(output_img)
+    output_img = output_img.resize(original_size, Image.BICUBIC)
+    output_img = np.array(output_img)
+
+    # Sharpen image (balanced)
+    blurred = cv2.GaussianBlur(output_img, (0,0), 1)
+    output_img = cv2.addWeighted(output_img, 1.3, blurred, -0.3, 0)
+
+    # Clip values
+    output_img = np.clip(output_img, 0, 255)
+
+    # Save for download
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
     Image.fromarray(output_img).save(temp_file.name)
 
     return output_img, temp_file.name
 
+# ---------------------------
+# GRADIO UI
+# ---------------------------
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔍 Image Super Resolution")
+    gr.Markdown("# 🔍 AI Image Enhancer")
+    gr.Markdown("Upload a low-quality image and enhance it using deep learning")
+
+    with gr.Row():
+        input_img = gr.Image(type="pil", label="Upload Image")
+        output_img = gr.Image(label="Enhanced Image")
 
-    input_img = gr.Image(type="pil", label="Upload Image")
-    output_img = gr.Image(label="Enhanced Image")
-    download_file = gr.File(label="Download Image")
+    download_file = gr.File(label="Download Enhanced Image")
 
     btn = gr.Button("Enhance Image")
 
@@ -106,19 +122,4 @@ with gr.Blocks() as demo:
         outputs=[output_img, download_file]
     )
 
-demo.launch()
-
-# ---------------------------
-# GRADIO UI
-# ---------------------------
-
-interface = gr.Interface(
-    fn=enhance_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="numpy", label="Enhanced Image"),
-    title="🔍 Super Resolution App",
-    description="Upload a low-quality image and enhance it using deep learning",
-    allow_flagging="never"
-)
-
-interface.launch()
+demo.launch()