app.py

import gradio as gr
import cv2
import numpy as np
from sklearn.cluster import KMeans
from PIL import Image
import math

# --- 1. 建立多組噴漆色卡資料庫 ---
# 這些 RGB 數值是根據官方色卡與常見數值進行的模擬近似值

COLOR_PALETTES = {
    "Generic (通用標準色)": {
        "Black": (0, 0, 0),
        "White": (255, 255, 255),
        "Dark Grey": (64, 64, 64),
        "Medium Grey": (128, 128, 128),
        "Light Grey": (192, 192, 192),
        "Red": (255, 0, 0),
        "Green": (0, 255, 0),
        "Blue": (0, 0, 255),
        "Yellow": (255, 255, 0),
        "Orange": (255, 165, 0),
        "Purple": (128, 0, 128),
        "Pink": (255, 192, 203),
        "Cyan": (0, 255, 255),
        "Magenta": (255, 0, 255),
        "Brown": (165, 42, 42)
    },
    "Montana Black (高壓/街頭風格)": {
        "BLK-9001 Black": (0, 0, 0),
        "BLK-9105 White": (255, 255, 255),
        "BLK-2093 Code Red": (227, 6, 19),
        "BLK-3145 Pimp Violet": (89, 36, 124),
        "BLK-4050 Royal Blue": (28, 62, 147),
        "BLK-4230 Tiffany": (0, 167, 157),
        "BLK-5030 Kicking Yellow": (255, 222, 0),
        "BLK-6035 Irish Green": (0, 142, 58),
        "BLK-6160 Ms. Piggy": (245, 184, 197),
        "BLK-7080 Pan": (118, 184, 42),
        "BLK-7210 Ghetto": (123, 117, 110),
        "BLK-8020 Shark": (60, 60, 60),
        "BLK-8110 Silverchrome": (180, 180, 180),
        "BLK-1025 Melon": (253, 174, 53),
        "BLK-4010 Blue Lagoon": (118, 205, 216)
    },
    "Montana Gold (低壓/細緻創作)": {
        "G9001 Shock Black": (25, 25, 25), # Gold 的黑通常稍微軟一點
        "G9100 Shock White": (250, 250, 250),
        "S3000 Shock Red": (226, 0, 37),
        "S5000 Shock Blue": (0, 107, 182),
        "S1000 Shock Yellow": (255, 237, 0),
        "G6130 Malachite": (0, 163, 120),
        "G2070 Orange": (255, 108, 47),
        "G3130 Pink Pink": (236, 0, 140),
        "G4170 Blue Velvet": (73, 75, 153),
        "G1210 Yellow Submarine": (255, 204, 0),
        "G6050 Greenery": (132, 189, 0),
        "G7070 Gravel": (125, 127, 124),
        "G5050 Iron Curtain": (88, 93, 104),
        "M1000 Gold Chrome": (212, 175, 55),
        "G3020 Strawberry": (234, 59, 78)
    },
    "Molotow Premium (色彩飽和)": {
        "221 Deep Black": (10, 10, 10),
        "231 Signal White": (255, 255, 255),
        "013 Dare Orange": (255, 88, 0),
        "086 Burgundy": (120, 20, 40),
        "027 Loomit Apricot": (246, 173, 131),
        "214 Signal Black": (20, 20, 20),
        "212 Stone Grey Middle": (130, 132, 133),
        "004 Signal Yellow": (255, 205, 0),
        "150 Kiwi": (140, 195, 60),
        "114 Riviera Middle": (0, 150, 200),
        "127 Sapphire Blue": (30, 60, 120),
        "069 Currant": (110, 40, 110),
        "058 Fuchsia Pink": (225, 25, 120),
        "096 Tulip Blue Middle": (80, 170, 210),
        "169 Poison Green": (0, 130, 60)
    }
}

def get_closest_color_name(rgb_tuple, palette_name):
    """
    根據選定的色表 (palette_name)，找出與輸入 RGB 最接近的噴漆顏色
    """
    # 取得對應的字典，如果找不到則預設用 Generic
    current_palette = COLOR_PALETTES.get(palette_name, COLOR_PALETTES["Generic (通用標準色)"])
    
    min_dist = float('inf')
    closest_name = "Unknown"
    
    r1, g1, b1 = rgb_tuple
    
    for name, (r2, g2, b2) in current_palette.items():
        # 計算歐幾里得距離
        dist = math.sqrt((r1 - r2)**2 + (g1 - g2)**2 + (b1 - b2)**2)
        if dist < min_dist:
            min_dist = dist
            closest_name = name
            
    return closest_name

def rgb_to_hex(rgb):
    return '#{:02x}{:02x}{:02x}'.format(int(rgb[0]), int(rgb[1]), int(rgb[2]))

def resize_image(image, max_width=800):
    h, w = image.shape[:2]
    if w > max_width:
        ratio = max_width / w
        new_h = int(h * ratio)
        return cv2.resize(image, (max_width, new_h), interpolation=cv2.INTER_AREA)
    return image

def sort_centers_by_luminance(centers):
    # 計算亮度: 0.299R + 0.587G + 0.114B (人眼對綠色最敏感)
    luminance = np.sum(centers * np.array([0.299, 0.587, 0.114]), axis=1)
    return centers[np.argsort(luminance)]

def stencil_art_generator(input_img, n_colors, smoothness, palette_select):
    if input_img is None:
        return None, [], []
    
    img = np.array(input_img)
    img = resize_image(img)
    
    # 圖像平滑化
    sigma = smoothness * 10
    img_blur = cv2.bilateralFilter(img, 9, sigma, sigma)
    
    pixel_values = img_blur.reshape((-1, 3))
    pixel_values = np.float32(pixel_values)
    
    # K-Means 聚類
    kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
    labels = kmeans.fit_predict(pixel_values)
    centers = kmeans.cluster_centers_
    
    # 依照亮度排序
    centers = sort_centers_by_luminance(centers)
    
    # 重新對應 labels
    from scipy.spatial.distance import cdist
    distances = cdist(pixel_values, centers)
    labels = np.argmin(distances, axis=1)
    
    centers = np.uint8(centers)
    
    layers = []
    color_info_list = []
    
    h, w = img.shape[:2]
    final_composite = np.zeros((h, w, 3), dtype=np.uint8)
    final_composite[:] = 255 
    
    for i in range(n_colors):
        layer_img = np.zeros((h, w, 4), dtype=np.uint8)
        mask = labels.reshape(h, w) == i
        color = centers[i]
        
        # --- 根據選擇的色表找出最接近的顏色 ---
        hex_code = rgb_to_hex(color)
        paint_name = get_closest_color_name(color, palette_select)
        
        label_str = f"L{i+1}: {paint_name}"
        
        color_info_list.append([
            f"Layer {i+1}",      
            paint_name,          
            hex_code,            
            f"{color[0]}, {color[1]}, {color[2]}" 
        ])
        
        layer_img[mask] = [color[0], color[1], color[2], 255]
        layers.append((Image.fromarray(layer_img), label_str))
        
        final_composite[mask] = color

    return Image.fromarray(final_composite), layers, color_info_list

# --- Gradio 介面設定 ---

with gr.Blocks(title="Stencil Art Generator with Color Match") as demo:
    gr.Markdown("# 🎨 AI 噴漆模板產生器 (含品牌色號)")
    gr.Markdown("上傳圖片，AI 會解析圖層並建議所選品牌的對應噴漆顏色。")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="上傳圖片")
            
            with gr.Accordion("參數設定", open=True):
                # 新增：品牌色卡選擇選單
                palette_dropdown = gr.Dropdown(
                    choices=list(COLOR_PALETTES.keys()),
                    value="Generic (通用標準色)",
                    label="選擇噴漆品牌 / 色卡"
                )
                slider_colors = gr.Slider(2, 10, value=5, step=1, label="圖層數量")
                slider_smooth = gr.Slider(1, 10, value=3, step=1, label="平滑度")
            
            btn_run = gr.Button("開始分析", variant="primary")
        
        with gr.Column(scale=2):
            output_final = gr.Image(type="pil", label="最終合成預覽")
            
            gr.Markdown("### 📋 建議色號清單")
            output_table = gr.Dataframe(
                headers=["圖層", "建議噴漆顏色 (品牌)", "Hex 色碼", "RGB 數值"],
                datatype=["str", "str", "str", "str"],
                label="Color Palette"
            )
            
            gr.Markdown("### 🖼️ 獨立圖層 (可下載)")
            output_gallery = gr.Gallery(label="圖層預覽", columns=3, height="auto")

    btn_run.click(
        fn=stencil_art_generator, 
        inputs=[input_image, slider_colors, slider_smooth, palette_dropdown], 
        outputs=[output_final, output_gallery, output_table]
    )

if __name__ == "__main__":
    demo.launch()