app.py

import cv2
import numpy as np
from collections import defaultdict
from scipy import stats
import gradio as gr

# ----------------------------
# Phân cụm histogram tự động
# ----------------------------

def update_c(C, hist):
    while True:
        groups = defaultdict(list)
        for i in range(len(hist)):
            if hist[i] == 0:
                continue
            d = np.abs(C - i)
            index = np.argmin(d)
            groups[index].append(i)

        new_C = np.array(C)
        for i, indice in groups.items():
            if np.sum(hist[indice]) == 0:
                continue
            new_C[i] = int(np.sum(indice * hist[indice]) / np.sum(hist[indice]))

        if np.sum(new_C - C) == 0:
            break
        C = new_C

    return C, groups

def K_histogram(hist):
    alpha = 0.001
    N = 80
    C = np.array([128])

    while True:
        C, groups = update_c(C, hist)
        new_C = set()
        for i, indice in groups.items():
            if len(indice) < N:
                new_C.add(C[i])
                continue

            z, pval = stats.normaltest(hist[indice])
            if pval < alpha:
                left = 0 if i == 0 else C[i - 1]
                right = len(hist) - 1 if i == len(C) - 1 else C[i + 1]
                delta = right - left
                if delta >= 3:
                    c1 = (C[i] + left) / 2
                    c2 = (C[i] + right) / 2
                    new_C.add(c1)
                    new_C.add(c2)
                else:
                    new_C.add(C[i])
            else:
                new_C.add(C[i])

        if len(new_C) == len(C):
            break
        else:
            C = np.array(sorted(new_C))
    return C

# ----------------------------
# Hàm cartoon hóa ảnh
# ----------------------------

def caart(img, bilateral_d=5, bilateral_sigmaColor=150, bilateral_sigmaSpace=150,
          canny_thresh1=100, canny_thresh2=200, erosion_iter=1):
    kernel = np.ones((2, 2), np.uint8)
    output = np.array(img)
    x, y, c = output.shape

    for i in range(c):
        output[:, :, i] = cv2.bilateralFilter(output[:, :, i], bilateral_d, bilateral_sigmaColor, bilateral_sigmaSpace)

    edge = cv2.Canny(output, canny_thresh1, canny_thresh2)
    output = cv2.cvtColor(output, cv2.COLOR_RGB2HSV)

    hists = []
    for i in range(c):
        bins = np.arange(180 + 1) if i == 0 else np.arange(256 + 1)
        hist, _ = np.histogram(output[:, :, i], bins=bins)
        hists.append(hist)

    C = [K_histogram(h) for h in hists]

    output = output.reshape((-1, c))
    for i in range(c):
        channel = output[:, i]
        index = np.argmin(np.abs(channel[:, np.newaxis] - C[i]), axis=1)
        output[:, i] = C[i][index]
    output = output.reshape((x, y, c))
    output = cv2.cvtColor(output, cv2.COLOR_HSV2RGB)

    contours, _ = cv2.findContours(edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    cv2.drawContours(output, contours, -1, 0, thickness=1)

    for i in range(3):
        output[:, :, i] = cv2.erode(output[:, :, i], kernel, iterations=erosion_iter)

    return output

# ----------------------------
# Hàm xử lý Gradio
# ----------------------------

def process(img, bilateral_d, bilateral_sigmaColor, bilateral_sigmaSpace, canny_thresh1, canny_thresh2, erosion_iter):
    return caart(img, bilateral_d, bilateral_sigmaColor, bilateral_sigmaSpace, canny_thresh1, canny_thresh2, erosion_iter)

# ----------------------------
# Preset cấu hình hiệu ứng
# ----------------------------

presets = {
    "Mặc định": [5, 150, 150, 100, 200, 1],
    "Phong cảnh": [9, 200, 200, 80, 160, 1],
    "Chân dung": [3, 100, 100, 120, 250, 0],
    "Truyện tranh / Manga": [5, 100, 100, 50, 100, 3],
    "Màu nước": [12, 250, 250, 100, 200, 1],
    "Tĩnh vật / đồ vật": [5, 150, 150, 80, 180, 1],
    "Ảnh thiếu sáng": [5, 100, 100, 60, 140, 0],
    "Sketch / thô mộc": [5, 80, 80, 50, 110, 4],
    "Trừu tượng / nghệ thuật": [9, 200, 200, 100, 250, 2],
}

def update_from_preset(preset_name):
    return presets.get(preset_name, presets["Mặc định"])

# ----------------------------
# Giao diện Gradio
# ----------------------------

with gr.Blocks() as demo:
    gr.Markdown("🎨 **Cartoonizer App** – Chuyển ảnh thật thành tranh hoạt hình bằng phân cụm màu và viền nét tay.")

    with gr.Row():
        img_input = gr.Image(type="numpy", label="Ảnh gốc")

        with gr.Column():
            preset = gr.Dropdown(
                choices=list(presets.keys()),
                value="Mặc định",
                label="🎛️ Chọn phong cách (Preset)"
            )

            bilateral_d = gr.Slider(1, 15, step=1, value=5, label="Bilateral Diameter")
            sigma_color = gr.Slider(50, 250, step=10, value=150, label="Sigma Color")
            sigma_space = gr.Slider(50, 250, step=10, value=150, label="Sigma Space")
            canny1 = gr.Slider(50, 200, step=10, value=100, label="Ngưỡng Canny thấp")
            canny2 = gr.Slider(100, 300, step=10, value=200, label="Ngưỡng Canny cao")
            erosion = gr.Slider(0, 5, step=1, value=1, label="Số lần Erode")

            preset.change(
                fn=update_from_preset,
                inputs=[preset],
                outputs=[bilateral_d, sigma_color, sigma_space, canny1, canny2, erosion]
            )

    cartoon_output = gr.Image(type="numpy", label="Ảnh hoạt hình")

    run_btn = gr.Button("🖌️ Tạo ảnh hoạt hình")
    run_btn.click(
        fn=process,
        inputs=[img_input, bilateral_d, sigma_color, sigma_space, canny1, canny2, erosion],
        outputs=cartoon_output
    )

demo.launch()