Update app.py

app.py

@@ -1,21 +1,102 @@
-# -----------------------------------------
-# DOWNLOAD MODEL FILES IF MISSING
-# -----------------------------------------
-if not os.path.exists(PROTOTXT):
-    urllib.request.urlretrieve(
-        "https://raw.githubusercontent.com/richzhang/colorization/master/colorization/models/colorization_deploy_v2.prototxt",
-        PROTOTXT
-    )
+import gradio as gr
+import numpy as np
+from PIL import Image
+from skimage import color as skcolor
+
+# Import Zhang full model (download logic is inside zhang_colorizer.py)
+from zhang_colorizer import zhang_colorize
+
+# -----------------------------
+# TEST MODE (subtle tint)
+# -----------------------------
+def test_mode_colorize(pil_img: Image.Image) -> Image.Image:
+    img = pil_img.convert("RGB")
+    arr = np.array(img).astype(np.float32) / 255.0
+    tint = np.array([1.02, 1.0, 0.98], dtype=np.float32)
+    arr = np.clip(arr * tint, 0.0, 1.0)
+    arr = (arr * 255).astype(np.uint8)
+    return Image.fromarray(arr)
+
+# -----------------------------
+# DEOLDIFY‑LITE (CPU SAFE)
+# -----------------------------
+def deoldify_lite_colorize(pil_img: Image.Image) -> Image.Image:
+    img = pil_img.convert("RGB")
+    arr = np.array(img).astype(np.float32) / 255.0
+
+    lab = skcolor.rgb2lab(arr)
+    L = lab[:, :, 0]
+    a = lab[:, :, 1]
+    b = lab[:, :, 2]
+
+    a *= 1.35
+    b *= 1.35
+    a += 2.0
+    b += 1.0
+
+    lab_out = np.stack([L, a, b], axis=-1)
+    rgb_out = skcolor.lab2rgb(lab_out)
+    rgb_out = np.clip(rgb_out, 0.0, 1.0)
+    rgb_out = (rgb_out * 255).astype(np.uint8)
+
+    return Image.fromarray(rgb_out)
+
+# -----------------------------
+# MAIN PIPELINE
+# -----------------------------
+def colorize_image(input_image, mode):
+    if input_image is None:
+        return None
 
-if not os.path.exists(CAFFEMODEL):
-    urllib.request.urlretrieve(
-        "https://raw.githubusercontent.com/richzhang/colorization/master/colorization/models/colorization_release_v2.caffemodel",
-        CAFFEMODEL
+    pil_img = input_image.convert("RGB")
+
+    if mode == "Test Mode (Very Subtle)":
+        return test_mode_colorize(pil_img)
+
+    if mode == "DeOldify‑Lite (Art Mode)":
+        return deoldify_lite_colorize(pil_img)
+
+    if mode == "Zhang ECCV (Realistic Strong Color)":
+        return zhang_colorize(pil_img)
+
+    return test_mode_colorize(pil_img)
+
+# -----------------------------
+# GRADIO UI
+# -----------------------------
+with gr.Blocks(title="Biker Image Colorizer – CPU (3 Modes)") as demo:
+    gr.Markdown(
+        """
+        # Biker Image Colorizer – CPU Edition  
+        **Three modes:**  
+        - Test Mode (very subtle)  
+        - DeOldify‑Lite (artistic strong color)  
+        - Zhang ECCV (realistic strong color)  
+        """
     )
 
-if not os.path.exists(PTS):
-    urllib.request.urlretrieve(
-        "https://raw.githubusercontent.com/richzhang/colorization/master/colorization/resources/pts_in_hull.npy",
-        PTS
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            mode = gr.Radio(
+                choices=[
+                    "Test Mode (Very Subtle)",
+                    "DeOldify‑Lite (Art Mode)",
+                    "Zhang ECCV (Realistic Strong Color)",
+                ],
+                value="Zhang ECCV (Realistic Strong Color)",
+                label="Colorization Mode",
+            )
+            run_btn = gr.Button("Colorize", variant="primary")
+
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Output Image")
+
+    run_btn.click(
+        fn=colorize_image,
+        inputs=[input_image, mode],
+        outputs=[output_image],
     )
- 
+
+if __name__ == "__main__":
+    demo.launch()