Update app.py

app.py

@@ -1,104 +1,46 @@
-import os
-import cv2
-import numpy as np
-import urllib.request
-import gradio as gr
-
-# ------------------ Working Model URLs from Hugging Face ------------------
-color_proto_url = "https://huggingface.co/datasets/amitverma-ai/old-photo-files/resolve/main/colorization_deploy_v2.prototxt"
-color_model_url = "https://huggingface.co/datasets/amitverma-ai/old-photo-files/resolve/main/colorization_release_v2.caffemodel"
-color_pts_url   = "https://huggingface.co/datasets/amitverma-ai/old-photo-files/resolve/main/pts_in_hull.npy"
-
-
-# ------------------ Utility Functions ------------------
-
-def download_if_missing(url, dest):
-    if not os.path.exists(dest):
-        print(f"Downloading {dest}...")
-        urllib.request.urlretrieve(url, dest)
-
-# ------------------ Setup Models ------------------
-
-model_dir = "models/colorization"
-os.makedirs(model_dir, exist_ok=True)
-
-proto_path = os.path.join(model_dir, "colorization_deploy_v2.prototxt")
-model_path = os.path.join(model_dir, "colorization_release_v2.caffemodel")
-pts_path = os.path.join(model_dir, "pts_in_hull.npy")
-
-download_if_missing(color_proto_url, proto_path)
-download_if_missing(color_model_url, model_path)
-download_if_missing(color_pts_url, pts_path)
-
-# Load model
-net = cv2.dnn.readNetFromCaffe(proto_path, model_path)
-pts = np.load(pts_path)
-
-class8 = net.getLayerId("class8_ab")
-conv8 = net.getLayerId("conv8_313_rh")
-pts = pts.transpose().reshape(2, 313, 1, 1)
-net.getLayer(class8).blobs = [pts.astype(np.float32)]
-net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
-
-# ------------------ Main Processing Function ------------------
-
-def restore_old_photo(image, face_enhance=True, colorize=True, scratch_remove=True):
-    try:
-        original = image.copy()
-
-        # Scratch removal using median blur (lightweight alternative)
-        if scratch_remove:
-            image = cv2.medianBlur(image, 3)
-
-        # Face enhancement: simulate sharpening with unsharp masking
-        if face_enhance:
-            blur = cv2.GaussianBlur(image, (0, 0), 3)
-            image = cv2.addWeighted(image, 1.5, blur, -0.5, 0)
-
-        # Colorization
-        if colorize:
-            h, w = image.shape[:2]
-            img_rgb = (image.astype("float32") / 255.0)
-            img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2LAB)
-            l_channel = img_lab[:, :, 0]
-
-            net_input = cv2.resize(l_channel, (224, 224))
-            net_input -= 50
-            net.setInput(cv2.dnn.blobFromImage(net_input))
-            ab_dec = net.forward()[0, :, :, :].transpose((1, 2, 0))
-            ab_dec_us = cv2.resize(ab_dec, (w, h))
-
-            lab_output = np.concatenate((l_channel[:, :, np.newaxis], ab_dec_us), axis=2)
-            bgr_output = cv2.cvtColor(lab_output.astype("float32"), cv2.COLOR_LAB2BGR)
-            bgr_output = np.clip(bgr_output * 255, 0, 255).astype("uint8")
-            image = bgr_output
-
-        return image
-
-    except Exception as e:
-        print(f"Error during restoration: {e}")
-        return original
-
-# ------------------ Gradio UI ------------------
-
-with gr.Blocks(title="AI Old Photo Restorer") as demo:
-    gr.Markdown("## 🧓🎨 AI Old Photo Restorer\nUpload old B/W or damaged photos and restore them with colorization, scratch removal, and face enhancement.")
-
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="📷 Upload Old Photo", type="numpy")
-            face_toggle = gr.Checkbox(label="👤 Face Enhancement", value=True)
-            colorize_toggle = gr.Checkbox(label="🎨 Colorization", value=True)
-            scratch_toggle = gr.Checkbox(label="🩹 Scratch Removal", value=True)
-            run_button = gr.Button("✨ Restore Photo", variant="primary")
-
-        with gr.Column():
-            output_image = gr.Image(label="🧼 Restored Photo")
-
-    run_button.click(
-        fn=restore_old_photo,
-        inputs=[input_image, face_toggle, colorize_toggle, scratch_toggle],
-        outputs=output_image
-    )
+import os, urllib.request, cv2, numpy as np, gradio as gr
+
+urls = {
+    "proto": "https://raw.githubusercontent.com/richzhang/colorization/master/colorization/models/colorization_deploy_v2.prototxt",
+    "model": "http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel",
+    "pts":   "https://raw.githubusercontent.com/richzhang/colorization/master/colorization/resources/pts_in_hull.npy"
+}
+os.makedirs("models", exist_ok=True)
+for name, url in urls.items():
+    path = f"models/{name}"
+    if not os.path.exists(path):
+        print(f"Downloading {name}...")
+        urllib.request.urlretrieve(url, path)
+
+net = cv2.dnn.readNetFromCaffe("models/proto", "models/model")
+pts = np.load("models/pts")
+for layer in ["class8_ab","conv8_313_rh"]:
+    net.getLayer(net.getLayerId(layer)).blobs = [pts.transpose().reshape(2,313,1,1).astype(np.float32)] if "class8" in layer else [np.full([1,313],2.606,dtype=np.float32)]
+
+def colorize(img):
+    h, w = img.shape[:2]
+    L = cv2.cvtColor(img.astype("float32")/255, cv2.COLOR_BGR2LAB)[:,:,0]
+    net.setInput(cv2.dnn.blobFromImage(cv2.resize(L,(224,224))-50))
+    ab = net.forward()[0].transpose((1,2,0))
+    ab = cv2.resize(ab,(w,h))
+    lab = np.concatenate([L[:,:,None],ab],axis=2)
+    return cv2.cvtColor(np.clip(lab,0,255).astype("uint8"), cv2.COLOR_LAB2BGR)
+
+def restore(img, face=True, scratch=True, color=True):
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    if scratch:
+        img = cv2.medianBlur(img,3)
+    if face:
+        blur = cv2.GaussianBlur(img,(0,0),3)
+        img = cv2.addWeighted(img,1.5,blur,-0.5,0)
+    if color:
+        img = colorize(img)
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+with gr.Blocks() as demo:
+    inp = gr.Image(type="numpy")
+    opts = [gr.Checkbox(label=l, value=True) for l in ["Face Enhancement","Scratch Removal","Colorization"]]
+    out = gr.Image()
+    gr.Button("Restore").click(restore, [inp]+opts, out)
 
 demo.launch()