#!/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()