Spaces:
Sleeping
Sleeping
File size: 2,941 Bytes
fe44dd3 2647161 fe44dd3 2647161 fe44dd3 bda4f37 fe44dd3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | 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
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)
"""
)
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()
|