Update app.py

app.py

@@ -3,12 +3,11 @@ import cv2
 import numpy as np
 import gradio as gr
 
-# ✅ Local model file paths (upload manually in Hugging Face Files tab)
+# Load pretrained models
 proto_path = "colorization_deploy_v2.prototxt"
 model_path = "colorization_release_v2.caffemodel"
 pts_path = "pts_in_hull.npy"
 
-# ✅ Load DNN colorization model
 net = cv2.dnn.readNetFromCaffe(proto_path, model_path)
 pts = np.load(pts_path)
 pts = pts.transpose().reshape(2, 313, 1, 1).astype(np.float32)
@@ -23,33 +22,41 @@ def colorize(img_bgr):
     L_rs = cv2.resize(L, (224, 224))
     L_rs -= 50
     net.setInput(cv2.dnn.blobFromImage(L_rs))
-    ab_dec = net.forward()[0, :, :, :].transpose((1, 2, 0))
+    ab_dec = net.forward()[0].transpose((1, 2, 0))
     ab_dec_us = cv2.resize(ab_dec, (w, h))
     lab_output = np.concatenate((L[:, :, np.newaxis], ab_dec_us), axis=2)
     bgr_output = cv2.cvtColor(lab_output.astype("uint8"), cv2.COLOR_LAB2BGR)
     return bgr_output
 
-def restore(img, face=True, scratch=True, colorize_img=True):
+def restore(img, enhance_face, remove_scratches, apply_colorization):
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    if scratch:
+    if remove_scratches:
         img = cv2.medianBlur(img, 3)
-    if face:
+    if enhance_face:
         blur = cv2.GaussianBlur(img, (0, 0), 3)
         img = cv2.addWeighted(img, 1.5, blur, -0.5, 0)
-    if colorize_img:
+    if apply_colorization:
         img = colorize(img)
-    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    restored = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    cv2.imwrite("restored_output.png", restored)
+    return restored, "restored_output.png"
+
+with gr.Blocks(title="AI Old Photo Restorer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🧓 AI Old Photo Restorer")
+    gr.Markdown("Give new life to your old, damaged, or black-and-white photos with facial enhancement, scratch removal, and colorization!")
 
-with gr.Blocks(title="AI Old Photo Restorer") as demo:
-    gr.Markdown("## 🧓📸 AI Old Photo Restorer")
     with gr.Row():
-        with gr.Column():
-            inp = gr.Image(label="Upload Photo", type="numpy")
-            face = gr.Checkbox("Face Enhancement", value=True)
-            scratch = gr.Checkbox("Scratch Removal", value=True)
-            color = gr.Checkbox("Colorization", value=True)
-            btn = gr.Button("Restore")
-        out = gr.Image(label="Restored Output")
-    btn.click(restore, [inp, face, scratch, color], out)
+        with gr.Column(scale=1):
+            inp = gr.Image(label="📤 Upload Old Photo", type="numpy", image_mode='RGB')
+            enhance_face = gr.Checkbox(label="✨ Enhance Facial Details", value=True)
+            remove_scratches = gr.Checkbox(label="🩹 Remove Scratches", value=True)
+            apply_colorization = gr.Checkbox(label="🎨 Colorize Black & White Photo", value=True)
+            btn = gr.Button("🔁 Restore Photo", variant="primary")
+            clear_btn = gr.ClearButton([inp, enhance_face, remove_scratches, apply_colorization])
+        with gr.Column(scale=1):
+            out_image = gr.Image(label="✅ Restored Output", interactive=False)
+            download = gr.File(label="⬇️ Download Restored Image")
+
+    btn.click(restore, [inp, enhance_face, remove_scratches, apply_colorization], [out_image, download])
 
 demo.launch()