Update app.py

app.py

@@ -1,128 +1,103 @@
-import gradio as gr
+import os
 import cv2
 import numpy as np
-import os
 import urllib.request
-import subprocess
-import psutil
-import tempfile
+import gradio as gr
 
-# ------------------ Config ------------------
-MAX_DIM = 1024  # Resize large images to avoid OOM
-device = 'cpu'
+# ------------------ Working Model URLs from Hugging Face ------------------
+color_proto_url = "https://huggingface.co/akhaliq/old-photo-restoration/resolve/main/colorization_deploy_v2.prototxt"
+color_model_url = "https://huggingface.co/akhaliq/old-photo-restoration/resolve/main/colorization_release_v2.caffemodel"
+color_pts_url = "https://huggingface.co/akhaliq/old-photo-restoration/resolve/main/pts_in_hull.npy"
 
-# ------------------ Model URLs ------------------
-color_proto_url = "https://github.com/richzhang/colorization/raw/caffe/models/colorization_deploy_v2.prototxt"
-color_model_url = "https://github.com/richzhang/colorization/raw/caffe/models/colorization_release_v2.caffemodel"
-color_pts_url = "https://github.com/richzhang/colorization/raw/caffe/resources/pts_in_hull.npy"
+# ------------------ Utility Functions ------------------
 
-# ------------------ Download Colorization Models ------------------
 def download_if_missing(url, dest):
     if not os.path.exists(dest):
         print(f"Downloading {dest}...")
         urllib.request.urlretrieve(url, dest)
 
-color_dir = "models/colorization"
-os.makedirs(color_dir, exist_ok=True)
-download_if_missing(color_proto_url, f"{color_dir}/colorization_deploy_v2.prototxt")
-download_if_missing(color_model_url, f"{color_dir}/colorization_release_v2.caffemodel")
-download_if_missing(color_pts_url, f"{color_dir}/pts_in_hull.npy")
-
-# ------------------ Load Colorization Net ------------------
-net = cv2.dnn.readNetFromCaffe(
-    f"{color_dir}/colorization_deploy_v2.prototxt",
-    f"{color_dir}/colorization_release_v2.caffemodel"
-)
-pts_in_hull = np.load(f"{color_dir}/pts_in_hull.npy")
+# ------------------ 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")
-net.getLayer(class8).blobs = [pts_in_hull.transpose().reshape(2, 313, 1, 1)]
-net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype=np.float32)]
-
-# ------------------ Optional GFPGAN lazy load ------------------
-gfpganer = None
-
-def enhance_face(image_np):
-    global gfpganer
-    if gfpganer is None:
-        from gfpgan import GFPGANer
-        gfpganer = GFPGANer(
-            model_path=None, upscale=1, arch='clean', channel_multiplier=2,
-            bg_upsampler=None, device=device
-        )
-    cropped_faces, _, restored_img = gfpganer.enhance(image_np, has_aligned=False, only_center_face=False, paste_back=True)
-    return restored_img
-
-# ------------------ Scratch Removal (OpenCV-based) ------------------
-def remove_scratches(img):
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    scratches = cv2.Laplacian(gray, cv2.CV_64F)
-    scratches = cv2.convertScaleAbs(scratches)
-    _, mask = cv2.threshold(scratches, 30, 255, cv2.THRESH_BINARY)
-    mask = cv2.dilate(mask, None, iterations=2)
-    inpainted = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
-    return inpainted
-
-# ------------------ Colorization ------------------
-def colorize_image(img):
-    h, w = img.shape[:2]
-    img_rs = cv2.resize(img, (224, 224))
-    lab = cv2.cvtColor(img_rs, cv2.COLOR_BGR2LAB)
-    l = lab[:, :, 0]
-    l -= 50
-
-    blob = cv2.dnn.blobFromImage(l)
-    net.setInput(blob)
-    ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
-    ab = cv2.resize(ab, (w, h))
-
-    L = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)[:, :, 0]
-    lab_out = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
-    bgr_out = cv2.cvtColor(lab_out.astype(np.uint8), cv2.COLOR_LAB2BGR)
-    return bgr_out
-
-# ------------------ Main Pipeline ------------------
-def restore_photo(image, enhance_face_flag, scratch_flag, color_flag):
-    print("Memory used (MB):", psutil.Process().memory_info().rss / 1024 / 1024)
-    
-    img = image
-    img_np = np.array(img)
-
-    # Resize if needed
-    h, w = img_np.shape[:2]
-    if max(h, w) > MAX_DIM:
-        scaling = MAX_DIM / max(h, w)
-        img_np = cv2.resize(img_np, (int(w * scaling), int(h * scaling)))
-
-    # Step-by-step processing
-    if scratch_flag:
-        img_np = remove_scratches(img_np)
-    if enhance_face_flag:
-        img_np = enhance_face(img_np)
-    if color_flag:
-        img_np = colorize_image(img_np)
-
-    return img_np
+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(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🧙‍♂️ AI Old Photo Restorer\nRestore, enhance, and colorize your old photos with AI on CPU!")
+
+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")
-            enhance_face = gr.Checkbox(label="🧑‍🦰 Enhance Faces", value=True)
-            remove_scratch = gr.Checkbox(label="🧽 Remove Scratches", value=True)
-            colorize = gr.Checkbox(label="🎨 Colorize Image", value=True)
-            run_btn = gr.Button("🚀 Restore")
+            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")
+            output_image = gr.Image(label="🧼 Restored Photo")
 
-    run_btn.click(
-        fn=restore_photo,
-        inputs=[input_image, enhance_face, remove_scratch, colorize],
+    run_button.click(
+        fn=restore_old_photo,
+        inputs=[input_image, face_toggle, colorize_toggle, scratch_toggle],
         outputs=output_image
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()