app.py

import gradio as gr
import numpy as np
from PIL import Image
import io
import torch
import torchvision.transforms as T

# -----------------------------
# GLOBAL FLAGS / DEVICE
# -----------------------------
DEVICE = "cpu"  # Force CPU for Hugging Face CPU Basic


# -----------------------------
# TEST MODE (your existing lightweight placeholder)
# -----------------------------
def test_mode_colorize(pil_img: Image.Image) -> Image.Image:
    """
    Very lightweight "test" colorizer.
    Just adds a tiny warm tint so we know the pipeline works.
    """
    img = pil_img.convert("RGB")
    arr = np.array(img).astype(np.float32) / 255.0

    # Subtle warm tint
    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)


# -----------------------------
# ZHANG ECCV16 COLORIZER (CPU)
# -----------------------------
# We use the lightweight PyTorch implementation from richzhang/colorization
# via the rz-colorization package (CPU‑friendly).
try:
    import colorizers  # from rz-colorization or richzhang/colorization
    _ZHANG_AVAILABLE = True
except Exception:
    _ZHANG_AVAILABLE = False
    colorizers = None

_zhang_model = None
_zhang_transform = T.Compose([
    T.Resize(256),
    T.CenterCrop(256),
    T.ToTensor(),
])


def load_zhang_model():
    global _zhang_model
    if not _ZHANG_AVAILABLE:
        return None
    if _zhang_model is None:
        # ECCV16 model = fully automatic colorization
        _zhang_model = colorizers.eccv16().eval().to(DEVICE)
    return _zhang_model


def zhang_colorize(pil_img: Image.Image) -> Image.Image:
    """
    Use Zhang ECCV16 model to colorize a grayscale image.
    Stronger, real colorization. CPU‑friendly.
    """
    model = load_zhang_model()
    if model is None:
        # Fallback: if model not available, just return test mode result
        return test_mode_colorize(pil_img)

    # Convert to RGB then to Lab
    img = pil_img.convert("RGB")
    img_resized = _zhang_transform(img)  # [3, H, W] in [0,1]
    img_resized = img_resized.unsqueeze(0).to(DEVICE)

    # Convert to Lab
    from skimage import color as skcolor

    np_img = img_resized[0].permute(1, 2, 0).cpu().numpy()
    lab = skcolor.rgb2lab(np_img)
    L = lab[:, :, 0]  # [H, W]

    tens_l = torch.from_numpy(L).unsqueeze(0).unsqueeze(0).float().to(DEVICE)

    with torch.no_grad():
        out_ab = model(tens_l).cpu()  # [1, 2, H, W]

    out_ab = out_ab[0].permute(1, 2, 0).numpy()  # [H, W, 2]

    # Resize ab to original image size
    H_orig, W_orig = img.size[1], img.size[0]
    out_ab_resized = np.array(
        Image.fromarray((out_ab * 255).astype(np.uint8)).resize((W_orig, H_orig), Image.BILINEAR),
        dtype=np.float32
    ) / 255.0

    # Rebuild Lab image at original resolution
    img_np = np.array(img).astype(np.float32) / 255.0
    lab_orig = skcolor.rgb2lab(img_np)
    L_orig = lab_orig[:, :, 0]

    lab_out = np.zeros((H_orig, W_orig, 3), dtype=np.float32)
    lab_out[:, :, 0] = L_orig
    lab_out[:, :, 1:] = out_ab_resized * 128.0  # scale back

    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)


# -----------------------------
# DEOLDIFY‑LITE (ART MODE, CPU)
# -----------------------------
# For CPU Basic, we implement a "lite" DeOldify‑style effect:
# stronger saturation + learned‑like color shift, but still lightweight.
def deoldify_lite_colorize(pil_img: Image.Image) -> Image.Image:
    """
    CPU‑safe "DeOldify‑Lite" style:
    - Convert to RGB
    - Apply stronger contrast + saturation
    - Apply learned‑like color mapping in Lab space
    """
    img = pil_img.convert("RGB")
    arr = np.array(img).astype(np.float32) / 255.0

    # Convert to Lab
    from skimage import color as skcolor

    lab = skcolor.rgb2lab(arr)
    L = lab[:, :, 0]
    a = lab[:, :, 1]
    b = lab[:, :, 2]

    # Boost chroma (a,b) to simulate stronger colorization
    a *= 1.35
    b *= 1.35

    # Gentle bias toward warmer tones
    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)":
        out = test_mode_colorize(pil_img)
    elif mode == "Zhang ECCV16 (Deep Colorizer)":
        out = zhang_colorize(pil_img)
    elif mode == "DeOldify‑Lite (Art Mode)":
        out = deoldify_lite_colorize(pil_img)
    else:
        out = test_mode_colorize(pil_img)

    return out


# -----------------------------
# GRADIO UI
# -----------------------------
with gr.Blocks(title="Fast CPU Biker Image Colorizer (Test + Zhang + DeOldify‑Lite)") as demo:
    gr.Markdown(
        """
        # Fast CPU Biker Image Colorizer (Test + Zhang + DeOldify‑Lite)

        **App 1 – Colorizer**

        - Runs entirely on **CPU Basic**
        - Three modes:
          - **Test Mode (Very Subtle)** – your original placeholder
          - **Zhang ECCV16 (Deep Colorizer)** – real neural colorization
          - **DeOldify‑Lite (Art Mode)** – stronger, art‑style color

        Upload a grayscale or sketch image on the left, choose a mode, then click **Colorize**.
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            mode = gr.Radio(
                choices=[
                    "Test Mode (Very Subtle)",
                    "Zhang ECCV16 (Deep Colorizer)",
                    "DeOldify‑Lite (Art Mode)",
                ],
                value="Zhang ECCV16 (Deep Colorizer)",
                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()