Agent/examples/find knowledge/scripts/extract_colors.py

#!/usr/bin/env python3
"""
提取色彩调色板 - 使用KMeans聚类
维度2: 全局色彩调色板 (color_palette)
维度7: 背景环境色彩 (background_color)
"""

import numpy as np
import json
import os
from PIL import Image, ImageDraw
from sklearn.cluster import KMeans
import colorsys

def rgb_to_hsl(r, g, b):
    """RGB转HSL"""
    r, g, b = r/255.0, g/255.0, b/255.0
    h, l, s = colorsys.rgb_to_hls(r, g, b)
    return {
        "h": round(h * 360, 1),
        "s": round(s * 100, 1),
        "l": round(l * 100, 1)
    }

def rgb_to_hex(r, g, b):
    return f"#{int(r):02x}{int(g):02x}{int(b):02x}"

def extract_palette(pixels, n_colors=6, img_id="", label="global"):
    """从像素数组提取主色调"""
    # 降采样加速
    if len(pixels) > 10000:
        idx = np.random.choice(len(pixels), 10000, replace=False)
        pixels_sample = pixels[idx]
    else:
        pixels_sample = pixels
    
    kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
    kmeans.fit(pixels_sample)
    
    # 计算每个聚类的比例
    labels = kmeans.predict(pixels_sample)
    unique, counts = np.unique(labels, return_counts=True)
    total = len(labels)
    
    colors = []
    for cluster_id, count in zip(unique, counts):
        center = kmeans.cluster_centers_[cluster_id]
        r, g, b = int(center[0]), int(center[1]), int(center[2])
        proportion = round(count / total, 3)
        
        colors.append({
            "rank": len(colors) + 1,
            "rgb": {"r": r, "g": g, "b": b},
            "hex": rgb_to_hex(r, g, b),
            "hsl": rgb_to_hsl(r, g, b),
            "proportion": proportion
        })
    
    # 按比例排序
    colors.sort(key=lambda x: x["proportion"], reverse=True)
    for i, c in enumerate(colors):
        c["rank"] = i + 1
    
    return colors

def create_palette_image(colors, width=600, height=100, output_path=None):
    """创建色彩调色板可视化图"""
    img = Image.new('RGB', (width, height), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    
    x = 0
    for color in colors:
        r, g, b = color["rgb"]["r"], color["rgb"]["g"], color["rgb"]["b"]
        block_width = int(color["proportion"] * width)
        if block_width > 0:
            draw.rectangle([x, 0, x + block_width, height], fill=(r, g, b))
            x += block_width
    
    if output_path:
        img.save(output_path)
    return img

def get_background_pixels(img_array, threshold_top=0.3, threshold_bottom=0.7):
    """提取背景区域像素（上部和左侧区域，排除人物主体区域）"""
    h, w = img_array.shape[:2]
    
    # 取图片上部（天空/树木区域）和左侧（背景区域）
    # 基于制作表分析：背景主要在左侧和上方
    top_region = img_array[:int(h * 0.4), :, :]  # 上40%
    left_region = img_array[:, :int(w * 0.35), :]  # 左35%
    
    # 合并背景区域
    bg_pixels = np.vstack([
        top_region.reshape(-1, 3),
        left_region.reshape(-1, 3)
    ])
    
    return bg_pixels

def main():
    input_dir = "input"
    
    # 全局色彩调色板
    palette_dir = "output/features/color_palette"
    # 背景色彩
    bg_dir = "output/features/background_color"
    
    palette_mappings = []
    bg_mappings = []
    
    # 段落对应关系
    segment_map = {
        "img_1": {"global_seg": "段落1", "bg_seg": "段落1.3"},
        "img_2": {"global_seg": "段落2", "bg_seg": "段落2.3"},
        "img_3": {"global_seg": "段落3", "bg_seg": "段落3.3"},
        "img_4": {"global_seg": "段落4", "bg_seg": "段落4.3"},
        "img_5": {"global_seg": "段落5", "bg_seg": "段落5.3"},
        "img_6": {"global_seg": "段落6", "bg_seg": "段落6.3"},
        "img_7": {"global_seg": "段落7", "bg_seg": "段落7.3"},
        "img_8": {"global_seg": "段落8", "bg_seg": "段落8.3"},
        "img_9": {"global_seg": "段落9", "bg_seg": "段落9.2"},
    }
    
    for i in range(1, 10):
        img_id = f"img_{i}"
        image_path = os.path.join(input_dir, f"{img_id}.jpg")
        
        if not os.path.exists(image_path):
            continue
        
        print(f"处理 {img_id}...")
        
        img = Image.open(image_path).convert('RGB')
        img_array = np.array(img)
        all_pixels = img_array.reshape(-1, 3).astype(float)
        
        # === 维度2：全局色彩调色板 ===
        global_colors = extract_palette(all_pixels, n_colors=6, img_id=img_id, label="global")
        
        # 保存JSON
        global_json_path = os.path.join(palette_dir, f"{img_id}_color_palette.json")
        with open(global_json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "image_id": img_id,
                "type": "global_color_palette",
                "n_colors": 6,
                "colors": global_colors,
                "description": "全局主色调，按比例排序"
            }, f, ensure_ascii=False, indent=2)
        
        # 保存可视化图
        global_img_path = os.path.join(palette_dir, f"{img_id}_color_palette.png")
        create_palette_image(global_colors, output_path=global_img_path)
        
        print(f"  ✓ 全局调色板: {[c['hex'] for c in global_colors[:3]]}")
        
        seg_info = segment_map.get(img_id, {})
        palette_mappings.append({
            "file": f"{img_id}_color_palette.png",
            "json_file": f"{img_id}_color_palette.json",
            "source_image": f"input/{img_id}.jpg",
            "segment": seg_info.get("global_seg", f"段落{i}"),
            "category": "形式",
            "feature": "整体色彩调色板（白绿配色）",
            "highlight_cluster": "cluster_3",
            "top_colors": [c["hex"] for c in global_colors[:3]]
        })
        
        # === 维度7：背景环境色彩 ===
        bg_pixels = get_background_pixels(img_array)
        bg_colors = extract_palette(bg_pixels.astype(float), n_colors=4, img_id=img_id, label="background")
        
        # 保存JSON
        bg_json_path = os.path.join(bg_dir, f"{img_id}_background_color.json")
        with open(bg_json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "image_id": img_id,
                "type": "background_color_palette",
                "n_colors": 4,
                "colors": bg_colors,
                "extraction_region": "上40%区域 + 左35%区域",
                "description": "背景区域主色调（自然绿色调）"
            }, f, ensure_ascii=False, indent=2)
        
        # 保存可视化图
        bg_img_path = os.path.join(bg_dir, f"{img_id}_background_color.png")
        create_palette_image(bg_colors, width=400, output_path=bg_img_path)
        
        print(f"  ✓ 背景色彩: {[c['hex'] for c in bg_colors[:3]]}")
        
        bg_mappings.append({
            "file": f"{img_id}_background_color.png",
            "json_file": f"{img_id}_background_color.json",
            "source_image": f"input/{img_id}.jpg",
            "segment": seg_info.get("bg_seg", ""),
            "category": "实质",
            "feature": "自然背景色彩（草地/树木绿色调）",
            "element_id": "元素3",
            "highlight_cluster": "cluster_3",
            "top_colors": [c["hex"] for c in bg_colors[:3]]
        })
    
    # 保存mapping.json for color_palette
    palette_mapping = {
        "dimension": "color_palette",
        "description": "全局色彩调色板，使用KMeans聚类提取6个主色调",
        "tool": "scikit-learn KMeans",
        "format": {
            "palette_image": "PNG，色块按比例排列，宽600px高100px",
            "palette_json": "JSON，包含6个主色调的RGB/HEX/HSL值和比例"
        },
        "mappings": palette_mappings
    }
    with open(os.path.join(palette_dir, "mapping.json"), 'w', encoding='utf-8') as f:
        json.dump(palette_mapping, f, ensure_ascii=False, indent=2)
    
    # 保存mapping.json for background_color
    bg_mapping = {
        "dimension": "background_color",
        "description": "背景区域色彩调色板，提取图片上部和左侧区域的主色调",
        "tool": "scikit-learn KMeans",
        "format": {
            "bg_image": "PNG，色块按比例排列，宽400px高100px",
            "bg_json": "JSON，包含4个主色调的RGB/HEX/HSL值和比例"
        },
        "mappings": bg_mappings
    }
    with open(os.path.join(bg_dir, "mapping.json"), 'w', encoding='utf-8') as f:
        json.dump(bg_mapping, f, ensure_ascii=False, indent=2)
    
    print("\n✓ 色彩调色板提取完成")

if __name__ == "__main__":
    main()