app.py

import gradio as gr
from PIL import Image, ImageFilter, ImageOps
import cv2
import numpy as np
import os
from collections import defaultdict
from skimage.color import deltaE_ciede2000, rgb2lab
import zipfile

def DoG_filter(image, kernel_size, sigma, k_sigma, gamma):
    if kernel_size == 0:
        kernel_size = None  # OpenCVにカーネルサイズを自動で決定させる
    else:
        kernel_size = (kernel_size, kernel_size)
    g1 = cv2.GaussianBlur(image, kernel_size, sigma)
    g2 = cv2.GaussianBlur(image, kernel_size, sigma * k_sigma)
    return g1 - gamma * g2

def XDoG_filter(image, kernel_size, sigma, k_sigma, epsilon, phi, gamma):
    epsilon /= 255
    dog = DoG_filter(image, kernel_size, sigma, k_sigma, gamma)
    dog /= dog.max()
    e = 1 + np.tanh(phi * (dog - epsilon))
    e[e >= 1] = 1
    return (e * 255).astype('uint8')

def process_XDoG(image, kernel_size, sigma, k_sigma, epsilon, phi, gamma):
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    xdog_image = XDoG_filter(image, kernel_size, sigma, k_sigma, epsilon, phi, gamma)
    return Image.fromarray(xdog_image).convert('L')



def replace_color(image, color_1, blur_radius=2):
    data = np.array(image)
    original_shape = data.shape
    channels = original_shape[2] if len(original_shape) > 2 else 1  # チャンネル数を確認
    data = data.reshape(-1, channels)
    color_1 = np.array(color_1)
    matches = np.all(data[:, :3] == color_1, axis=1)
    nochange_count = 0
    mask = np.zeros(data.shape[0], dtype=bool)

    while np.any(matches):
        new_matches = np.zeros_like(matches)
        match_num = np.sum(matches)
        for i in range(len(data)):  # Removed tqdm
            if matches[i]:
                x, y = divmod(i, original_shape[1])
                neighbors = [
                    (x, y-1), (x, y+1), (x-1, y), (x+1, y)
                ]
                valid_neighbors = []
                for nx, ny in neighbors:
                    if 0 <= nx < original_shape[0] and 0 <= ny < original_shape[1]:
                        ni = nx * original_shape[1] + ny
                        if not np.all(data[ni, :3] == color_1, axis=0):
                            valid_neighbors.append(data[ni, :3])
                if valid_neighbors:
                    new_color = np.mean(valid_neighbors, axis=0).astype(np.uint8)
                    data[i, :3] = new_color
                    data[i, 3] = 255
                    mask[i] = True
                else:
                    new_matches[i] = True
        matches = new_matches
        if match_num == np.sum(matches):
            nochange_count += 1
        if nochange_count > 5:
            break

    data = data.reshape(original_shape)
    mask = mask.reshape(original_shape[:2])

    result_image = Image.fromarray(data, 'RGBA')
    blurred_image = result_image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    blurred_data = np.array(blurred_image)

    np.copyto(data, blurred_data, where=mask[..., None])

    return Image.fromarray(data, 'RGBA')

def generate_distant_colors(consolidated_colors, distance_threshold):
    consolidated_lab = [rgb2lab(np.array([color], dtype=np.float32) / 255.0).reshape(3) for color, _ in consolidated_colors]
    max_attempts = 10000
    for _ in range(max_attempts):
        random_rgb = np.random.randint(0, 256, size=3)
        random_lab = rgb2lab(np.array([random_rgb], dtype=np.float32) / 255.0).reshape(3)
        if all(deltaE_ciede2000(base_color_lab, random_lab) > distance_threshold for base_color_lab in consolidated_lab):
            return tuple(random_rgb)
    return (128, 128, 128)

def consolidate_colors(major_colors, threshold):
    colors_lab = [rgb2lab(np.array([[color]], dtype=np.float32)/255.0).reshape(3) for color, _ in major_colors]
    i = 0
    while i < len(colors_lab):
        j = i + 1
        while j < len(colors_lab):
            if deltaE_ciede2000(colors_lab[i], colors_lab[j]) < threshold:
                if major_colors[i][1] >= major_colors[j][1]:
                    major_colors[i] = (major_colors[i][0], major_colors[i][1] + major_colors[j][1])
                    major_colors.pop(j)
                    colors_lab.pop(j)
                else:
                    major_colors[j] = (major_colors[j][0], major_colors[j][1] + major_colors[i][1])
                    major_colors.pop(i)
                    colors_lab.pop(i)
                continue
            j += 1
        i += 1
    return major_colors

def get_major_colors(image, threshold_percentage=0.01):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    color_count = defaultdict(int)
    for pixel in image.getdata():
        color_count[pixel] += 1
    total_pixels = image.width * image.height
    major_colors = [(color, count) for color, count in color_count.items() if (count / total_pixels) >= threshold_percentage]
    return major_colors

def line_color(image, mask, new_color):
    data = np.array(image)
    data[mask, :3] = new_color
    return Image.fromarray(data)

def process_image(image, lineart):
    if image.mode != 'RGBA':
        image = image.convert('RGBA')

    lineart = lineart.point(lambda x: 0 if x < 200 else 255)
    lineart = ImageOps.invert(lineart)
    kernel = np.ones((3, 3), np.uint8)
    lineart = cv2.dilate(np.array(lineart), kernel, iterations=1)
    lineart = Image.fromarray(lineart)
    mask = np.array(lineart) == 255
    major_colors = get_major_colors(image, threshold_percentage=0.05)
    major_colors = consolidate_colors(major_colors, 10)
    new_color_1 = generate_distant_colors(major_colors, 100)
    filled_image = line_color(image, mask, new_color_1)
    replace_color_image = replace_color(filled_image, new_color_1, 2).convert('RGB')
    return replace_color_image

def zip_files(zip_files, zip_path):
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file_path in zip_files:
            zipf.write(file_path, arcname=os.path.basename(file_path))

class webui:
    def __init__(self):
        self.demo = gr.Blocks()

    def main(self, image_path, kernel_size, sigma, k_sigma, epsilon, phi, gamma):
        kernel_size = int(kernel_size)
        sigma = float(sigma)
        k_sigma = float(k_sigma)
        epsilon = int(epsilon)
        phi = int(phi)
        gamma = float(gamma)
        
        image = Image.open(image_path).convert('RGBA')
        #拡張子を取り除いたファイル名を取得
        image_name = os.path.splitext(image_path)[0]
        alpha = image.getchannel('A') if image.mode == 'RGBA' else None
        image.save(image_path)
        image = Image.open(image_path).convert('RGBA')
        rgb_image = image.convert('RGB')
        lineart = process_XDoG(image, kernel_size, sigma, k_sigma, epsilon, phi, gamma).convert('L')
        replace_color_image = process_image(rgb_image, lineart).convert('RGBA')
        
        if alpha:
            replace_color_image.putalpha(alpha)
        
        replace_color_image_path = f"{image_name}_noline.png"
        replace_color_image.save(replace_color_image_path)
        
        lineart_image = lineart.convert('RGBA')
        lineart_alpha = 255 - np.array(lineart) 
        lineart_image.putalpha(Image.fromarray(lineart_alpha))
        
        lineart_image_path = f"{image_name}_lineart.png"
        lineart_image.save(lineart_image_path)

        zip_files_list = [replace_color_image_path, lineart_image_path]
        zip_path = f"{image_name}.zip"
        zip_files(zip_files_list, zip_path)
        
        outputs = [replace_color_image, lineart_image]
        return outputs, zip_path

    def launch(self, share):
        with self.demo:
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type='filepath', image_mode="RGBA", label="Original Image")
                    kernel_size = gr.Dropdown(choices=[0, 1, 3, 5, 7, 9], value=0, label="カーネルサイズ (kernel_size)")
                    kernel_size_description = gr.Label("ガウシアンフィルタのカーネルサイズです。線の鋭さを設定します。0を選択するとOpenCVが自動でカーネルサイズを決定します。")
                    sigma = gr.Slider(0.1, 10.0, step=0.1, value=1.4, label="シグマ (sigma)")
                    sigma_description = gr.Label("ガウシアンフィルタの標準偏差。ぼかしの程度を制御します。")
                    k_sigma = gr.Slider(1.0, 3.0, step=0.1, value=1.6, label="k_シグマ (k_sigma)")
                    k_sigma_description = gr.Label("二つ目のガウシアンフィルタのシグマの倍率。線の太さが影響を受けます。")
                    epsilon = gr.Slider(-10, 10, value=0, label="イプシロン (epsilon)")
                    epsilon_description = gr.Label("XDoGフィルターの閾値パラメータ。線画の明瞭度やコントラストを調整します。")
                    phi = gr.Slider(1, 100, value=10, label="ファイ (phi)")
                    phi_description = gr.Label("タンジェントハイパーボリック関数の係数。閾値を超えた部分の鮮明さを高めます。")
                    gamma = gr.Slider(0.5, 1.5, step=0.1, value=0.98, label="ガンマ (gamma)")
                    gamma_description = gr.Label("DoGフィルタの強度調整パラメータ。主に線の明るさや暗さを調節します。")
                    submit = gr.Button(value="Start")
                with gr.Column():
                    output_image = gr.Image(label="Processed Image")

            submit.click(
                self.main, 
                inputs=[input_image, kernel_size, sigma, k_sigma, epsilon, phi, gamma],
                outputs=[output_image]
            )

        self.demo.launch(share=share)

if __name__ == "__main__":
    ui = webui()
    ui.launch(share=True)