Update app.py

app.py

@@ -1,81 +1,128 @@
-import os, urllib.request
 import gradio as gr
-import numpy as np
 import cv2
-from PIL import Image
-import torch
-from gfpgan import GFPGANer
-
-# Download models if needed
-def download_models():
-    os.makedirs("models", exist_ok=True)
-    files = {
-      "models/colorization_deploy_v2.prototxt":
-        "https://raw.githubusercontent.com/richzhang/colorization/master/models/colorization_deploy_v2.prototxt",
-      "models/colorization_release_v2.caffemodel":
-        "http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel",
-      "models/pts_in_hull.npy":
-        "https://raw.githubusercontent.com/junyanz/interactive-deep-colorization/master/data/color_bins/pts_in_hull.npy"
-    }
-    for path, url in files.items():
-        if not os.path.exists(path):
-            print(f"🔽 Downloading {os.path.basename(path)} …")
-            urllib.request.urlretrieve(url, path)
-
-download_models()
-
-# Load GFPGAN for face enhancement
-gfpganer = GFPGANer(model_path=None, upscale=2, arch='clean', channel_multiplier=2,
-                    bg_upsampler=None, device='cpu')
-
-# Setup colorization network
-net = cv2.dnn.readNetFromCaffe("models/colorization_deploy_v2.prototxt",
-                               "models/colorization_release_v2.caffemodel")
-pts = np.load("models/pts_in_hull.npy")
-pts = pts.transpose().reshape(2, 313, 1, 1).astype("float32")
-net.getLayer(net.getLayerId('class8_ab')).blobs = [pts]
-net.getLayer(net.getLayerId('conv8_313_rh')).blobs = [np.full((1,313),2.606,dtype="float32")]
-
-def detect_scratches(img):
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    _, mask = cv2.threshold(gray, 235, 255, cv2.THRESH_BINARY)
-    return cv2.dilate(mask, np.ones((3,3), np.uint8), iterations=1)
-
-def restore_photo(img, enhance_face, remove_scratch, colorize):
-    img_np = np.array(img.convert("RGB"))
-    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-
-    if remove_scratch:
-        mask = detect_scratches(img_np)
-        img_bgr = cv2.inpaint(img_bgr, mask, 3, cv2.INPAINT_TELEA)
-
-    if enhance_face:
-        _, img_bgr = gfpganer.enhance(img_bgr, has_aligned=False,
-                                      only_center_face=False, paste_back=True)
-
-    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
-
-    if colorize:
-        lab = cv2.cvtColor(img_rgb.astype("float32")/255.0, cv2.COLOR_RGB2Lab)
-        L = lab[:,:,0]
-        L_rs = cv2.resize(L, (224,224)) - 50
-        net.setInput(cv2.dnn.blobFromImage(L_rs))
-        ab = net.forward()[0].transpose((1,2,0))
-        ab_us = cv2.resize(ab, (img_rgb.shape[1], img_rgb.shape[0]))
-        lab_out = np.concatenate([L[:,:,None], ab_us], axis=2)
-        img_rgb = np.clip(cv2.cvtColor(lab_out, cv2.COLOR_Lab2RGB)*255, 0, 255).astype("uint8")
-
-    return Image.fromarray(img_rgb)
-
-with gr.Blocks(title="AI Photo Restorer") as demo:
-    gr.Markdown("### 🧓✨ AI Old Photo Restorer")
-    inp = gr.Image(label="Upload Photo", type="pil")
-    fe = gr.Checkbox("Enhance Faces", value=True)
-    sr = gr.Checkbox("Remove Scratches", value=True)
-    cz = gr.Checkbox("Colorize Image", value=True)
-    btn = gr.Button("Restore")
-    out = gr.Image(label="Restored Photo")
-    btn.click(restore_photo, inputs=[inp, fe, sr, cz], outputs=out)
-
-if __name__=="__main__":
+import numpy as np
+import os
+import urllib.request
+import subprocess
+import psutil
+import tempfile
+
+# ------------------ Config ------------------
+MAX_DIM = 1024  # Resize large images to avoid OOM
+device = 'cpu'
+
+# ------------------ 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"
+
+# ------------------ 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")
+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
+
+# ------------------ 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.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")
+        with gr.Column():
+            output_image = gr.Image(label="Restored Photo")
+
+    run_btn.click(
+        fn=restore_photo,
+        inputs=[input_image, enhance_face, remove_scratch, colorize],
+        outputs=output_image
+    )
+
+if __name__ == "__main__":
     demo.launch()