# -*- coding: utf-8 -*- """ ChartPipeline: 完整可视化生成管道 此脚本执行完整的数据可视化生成流程，依次调用各模块实现从数据到最终图表的转换 """ import os import json import argparse import logging import time from datetime import datetime from importlib import import_module from pathlib import Path from config import ( base_url, api_key, embed_model_path, data_resource_path, topk, text_data_path, text_index_path, color_data_path, color_index_path, image_data_path, image_index_path, image_list_path, image_resource_path ) import random from concurrent.futures import ProcessPoolExecutor # 配置日志 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger("ChartPipeline") # 模块配置 MODULES = [ { "name": "create_index", "description": "索引创建模块", "input_type": "none", "output_type": "none" }, { "name": "preprocess", "description": "数据预处理模块", "input_type": "json", "output_type": "json" }, { "name": "chart_type_recommender", "description": "图表类型推荐模块", "input_type": "json", "output_type": "json" }, { "name": "datafact_generator", "description": "数据洞察模块", "input_type": "json", "output_type": "json" }, { "name": "title_generator", "description": "标题生成模块", "input_type": "json", "output_type": "json" }, { "name": "color_recommender", "description": "色彩推荐模块", "input_type": "json", "output_type": "json" }, { "name": "image_recommender", "description": "图像推荐模块", "input_type": "json", "output_type": "json" }, { "name": "all", "description": "综合模块(datafact+title+color+image)", "input_type": "json", "output_type": "json" }, { "name": "infographics_generator", "description": "信息图表生成模块", "input_type": "json", "output_type": "json" }, { "name": "asset_regenerator", "description": "图像/图标统一生图后处理模块", "input_type": "svg", "output_type": "svg" }, { "name": "full_image_polisher", "description": "整图GPT Image润色后处理模块", "input_type": "png", "output_type": "png" }, { "name": "chart_engine", "description": "图表模板实现引擎", "input_type": "json", "output_type": "svg" }, { "name": "title_styler", "description": "标题元素生成模块", "input_type": "json", "output_type": "svg" } ] def run_pipeline(input_path, output_path=None, temp_dir=None, modules_to_run=None, threads=None, chart_name=None, chart_only=False, output_png=False): """ 执行完整的图表生成管道 Args: input_path (str): 输入数据文件路径(可以是文件或目录) output_path (str, optional): 输出文件路径(可以是文件或目录)，如果为None则原地修改 temp_dir (str, optional): 临时文件目录，默认使用./tmp modules_to_run (list, optional): 要运行的模块列表，默认运行所有模块 threads (int, optional): 处理目录时的并发线程数，仅在input_path为目录时生效 chart_name (str, optional): 指定图表名称，仅对infographics_generator模块有效 chart_only (bool, optional): 只输出 chart SVG（跳过 title/image/layout），仅对 infographics_generator 有效 output_png (bool, optional): 同时输出 PNG，仅对 infographics_generator 有效 """ # try: if chart_name is not None: output_path = os.path.join(output_path, chart_name) import shutil shutil.rmtree(output_path, ignore_errors=True) os.makedirs(output_path, exist_ok=True) print("modules_to_run: ", modules_to_run) # 如果是create_index模块，单独处理 if modules_to_run and 'create_index' in modules_to_run: return run_single_file( input_path=input_path, output_path=output_path, temp_dir=temp_dir, modules_to_run=modules_to_run, chart_name=chart_name, chart_only=chart_only, output_png=output_png, ) input_path = Path(input_path) output_path = Path(output_path) if output_path else input_path temp_dir = Path(temp_dir) if temp_dir else Path("./tmp") if input_path.is_dir(): # 如果指定了输出目录且不同于输入目录，创建输出目录 if output_path != input_path: output_path.mkdir(parents=True, exist_ok=True) # 获取输入目录下所有JSON文件 input_files = list(input_path.glob('*.json')) random.shuffle(input_files) # 随机打乱文件顺序 if threads and threads > 1: # 使用线程池并行处理文件 with ProcessPoolExecutor(max_workers=threads) as executor: futures = [] for input_file in input_files: # 如果是inplace处理，输出路径就是输入路径 if output_path == input_path: output_file = input_file else: # 确保输出文件保持相同的文件名 output_file = output_path / input_file.name future = executor.submit( run_single_file, input_path=input_file, output_path=output_file, temp_dir=temp_dir, modules_to_run=modules_to_run, chart_name=chart_name, chart_only=chart_only, output_png=output_png, ) futures.append(future) # 等待所有任务完成并检查结果 results = [future.result() for future in futures] return all(results) else: print("单线程顺序处理") # 单线程顺序处理 success = True for input_file in input_files: if output_path == input_path: output_file = input_file else: output_file = output_path / input_file.name success &= run_single_file( input_path=input_file, output_path=output_file, temp_dir=temp_dir, modules_to_run=modules_to_run, chart_name=chart_name, chart_only=chart_only, output_png=output_png, ) return success else: # 单文件处理 if output_path == input_path: output_file = input_path else: # 如果指定了不同的输出路径，确保它的父目录存在 output_path.parent.mkdir(parents=True, exist_ok=True) output_file = output_path return run_single_file( input_path=input_path, output_path=output_file, temp_dir=temp_dir, modules_to_run=modules_to_run, chart_name=chart_name, chart_only=chart_only, output_png=output_png, ) # except Exception as e: # logger.error(f"管道执行失败: {str(e)}") # return False def run_single_file(input_path, output_path, temp_dir=None, modules_to_run=None, chart_name=None, chart_only=False, output_png=False, return_timing=False): """ 处理单个文件的管道逻辑 Args: input_path (Path): 输入JSON文件路径 output_path (Path): 输出路径 temp_dir (Path): 临时文件目录，默认为./tmp modules_to_run (list): 要运行的模块列表 chart_name (str, optional): 指定图表名称，仅对infographics_generator模块有效 chart_only (bool, optional): 只输出 chart SVG，仅对 infographics_generator 有效 output_png (bool, optional): 同时输出 PNG，仅对 infographics_generator 有效 """ # try: run_start_time = time.time() module_timings = {} def finish(success: bool): total_seconds = time.time() - run_start_time logger.info( "run_single_file timing: total=%.3fs modules=%s", total_seconds, {k: round(v, 3) for k, v in module_timings.items()}, ) if return_timing: return { "success": success, "total_seconds": total_seconds, "module_seconds": module_timings, } return success def fail_module(module_name: str): logger.error("模块执行失败: %s", module_name) return finish(False) # 确定要运行的模块 if not modules_to_run: modules_to_run = [m["name"] for m in MODULES] # 如果要运行create_index，单独处理并直接返回 if 'create_index' in modules_to_run: logger.info("执行create_index模块") # 依次调用三个模块的create_index for module_type in ['title', 'color', 'image']: if module_type == 'title': print("Running modules.title_generator.create_index") module = import_module('modules.title_generator.create_index') if not Path(text_index_path).exists(): module.process( data=text_data_path, index_path=text_index_path, data_path=text_data_path, embed_model_path=embed_model_path ) elif module_type == 'color': print("Running modules.color_recommender.create_index") module = import_module('modules.color_recommender.create_index') if not Path(color_index_path).exists(): module.main( input=color_data_path, output=color_index_path, embed_model_path=embed_model_path ) else: # image print("Running modules.image_recommender.create_index") module = import_module('modules.image_recommender.create_index') if not Path(image_index_path).exists(): module.main( image_list_path=image_list_path, image_resource_path=image_resource_path, index_path=image_index_path, data_path=image_data_path, embed_model_path=embed_model_path ) total_seconds = time.time() - run_start_time if return_timing: return { "success": True, "total_seconds": total_seconds, "module_seconds": module_timings, } return True # 创建临时目录 if temp_dir is None: temp_dir = Path("./tmp") temp_dir.mkdir(parents=True, exist_ok=True) # 确保输出目录存在 output_path.parent.mkdir(parents=True, exist_ok=True) # 当前输入文件 current_input = input_path # 获取最后一个模块的配置 last_module = [m for m in MODULES if m["name"] in modules_to_run][-1] # 根据最后一个模块的输出类型决定最终输出文件的扩展名 if last_module["name"] in ["chart_engine", "title_styler", "asset_regenerator"]: final_output = output_path.with_suffix('.svg') elif last_module["name"] == "full_image_polisher": final_output = output_path.with_suffix('.png') elif last_module["name"] == "infographics_generator": final_output = output_path.parent / f"{output_path.stem}_final.svg" else: final_output = output_path.with_suffix('.json') # 记录执行过程 # 依次执行各模块 for i, module_config in enumerate([m for m in MODULES if m["name"] in modules_to_run]): module_name = module_config["name"] module_desc = module_config["description"] module_start_time = time.time() # logger.info(f"执行模块 {i+1}/{len(modules_to_run)}: {module_name} - {module_desc}") # 特殊处理all模块，依次执行datafact_generator、title_generator、color_recommender和image_recommender if module_name == "all": # 数据洞察模块 preprocess_module = import_module(f"modules.preprocess.preprocess") if not should_skip_module(module_name, output_path): step_start_time = time.time() preprocess_module.process(input=str(current_input), output=str(output_path)) module_timings["all.preprocess"] = time.time() - step_start_time current_input = output_path datafact_module = import_module("modules.datafact_generator.datafact_generator") if not should_skip_module("datafact_generator", output_path): step_start_time = time.time() datafact_module.process(input=str(current_input), output=str(output_path)) module_timings["all.datafact_generator"] = time.time() - step_start_time current_input = output_path # 标题生成模块 title_module = import_module("modules.title_generator.title_generator") if not should_skip_module("title_generator", output_path): step_start_time = time.time() title_module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, topk=topk, data_path=text_data_path, index_path=text_index_path ) module_timings["all.title_generator"] = time.time() - step_start_time current_input = output_path # 色彩推荐模块 color_module = import_module("modules.color_recommender.color_recommender") if not should_skip_module("color_recommender", output_path): step_start_time = time.time() color_module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, data_path=color_data_path, index_path=color_index_path ) module_timings["all.color_recommender"] = time.time() - step_start_time current_input = output_path # 图像推荐模块 image_module = import_module("modules.image_recommender.image_recommender") if not should_skip_module("image_recommender", output_path): step_start_time = time.time() image_module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, data_path=image_data_path, index_path=image_index_path, resource_path=image_resource_path ) module_timings["all.image_recommender"] = time.time() - step_start_time current_input = output_path # 特殊处理title_generator模块，传入配置参数 elif module_name == "title_generator": module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): step_start_time = time.time() module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, topk=topk, data_path=text_data_path, index_path=text_index_path ) module_timings["title_generator"] = time.time() - step_start_time current_input = output_path elif module_name == "color_recommender": module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): step_start_time = time.time() module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, data_path=color_data_path, index_path=color_index_path ) module_timings["color_recommender"] = time.time() - step_start_time current_input = output_path elif module_name == "image_recommender": module = import_module(f"modules.{module_name}.{module_name}") print("image_recommender") if not should_skip_module(module_name, output_path): print("process") step_start_time = time.time() module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, embed_model_path=embed_model_path, data_path=image_data_path, index_path=image_index_path, resource_path=image_resource_path ) module_timings["image_recommender"] = time.time() - step_start_time current_input = output_path elif module_name == "infographics_generator": module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): step_start_time = time.time() result = module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, chart_name=chart_name, chart_only=chart_only, output_png=output_png, ) module_timings["infographics_generator"] = time.time() - step_start_time if result is False: return fail_module(module_name) current_input = output_path elif module_name == "asset_regenerator": module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): ok = module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, ) if not ok: return False current_input = output_path elif module_name == "full_image_polisher": module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): ok = module.process( input=str(current_input), output=str(output_path), base_url=base_url, api_key=api_key, ) if not ok: return False current_input = output_path elif module_name == "chart_engine": # 输入是JSON，输出是SVG module = import_module(f"modules.{module_name}.{module_name}") svg_output = output_path.with_suffix('.svg') if True: # not svg_output.exists(): step_start_time = time.time() result = module.process(input=str(current_input), output=str(svg_output)) module_timings["chart_engine"] = time.time() - step_start_time if result is False: return fail_module(module_name) current_input = svg_output # 更新为SVG文件作为下一个模块的输入 elif module_name == "title_styler": # 输入是JSON，输出是SVG module = import_module(f"modules.{module_name}.{module_name}") title_svg = output_path.parent / f"{output_path.stem}_title.svg" if not title_svg.exists(): step_start_time = time.time() result = module.process(input=str(current_input), output=str(title_svg)) module_timings["title_styler"] = time.time() - step_start_time if result is False: return fail_module(module_name) current_input = title_svg # 更新为标题SVG作为下一个模块的输入 else: # 普通模块：输入JSON，输出JSON module = import_module(f"modules.{module_name}.{module_name}") if not should_skip_module(module_name, output_path): step_start_time = time.time() result = module.process(input=str(current_input), output=str(output_path)) module_timings[module_name] = time.time() - step_start_time if result is False: return fail_module(module_name) current_input = output_path module_timings[f"{module_name}.total"] = time.time() - module_start_time return finish(True) # except Exception as e: # logger.error(f"文件处理失败 {input_path}: {str(e)}") # return False def should_skip_module(module_name: str, output_path: Path) -> bool: """检查是否需要跳过模块执行""" try: if not output_path.exists(): return False # 检查JSON文件中的特定字段 with open(output_path) as f: data = json.load(f) skip_conditions = { "preprocess": lambda d: "metadata" in d and "data" in d and "variables" in d and "processed" in d, "chart_type_recommender": lambda d: "chart_type" in d, "datafact_generator": lambda d: "datafacts" in d, "title_generator": lambda d: False,#"titles" in d, "color_recommender": lambda d: "colors" in d, "image_recommender": lambda d: "images" in d } if module_name in skip_conditions: return skip_conditions[module_name](data) return False except Exception as e: logger.warning(f"检查跳过条件时出错: {str(e)}") return False def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--input', type=str, help='Input json file path or directory (default: config.data_resource_path = data_resource_dirs[0])', default=data_resource_path) parser.add_argument('--output', type=str, help='Output json file path', default='output') parser.add_argument('--temp-dir', type=str, default='tmp') parser.add_argument('--modules', type=str, nargs='+', help='Modules to run', default=['infographics_generator']) parser.add_argument('--threads', type=int, help='Number of threads for directory processing', default=1) parser.add_argument('--chart-name', type=str, help='Specific chart name to use for infographics_generator') parser.add_argument('--chart-only', action='store_true', help='Only output the chart SVG; skip title/image/layout (infographics_generator only)') parser.add_argument('--output-png', action='store_true', help='Also write a PNG next to the final SVG (uses rsvg-convert)') args = parser.parse_args() # 如果没有指定input，从data_resource_path随机选择 if args.input is None and 'create_index' not in args.modules: json_files = [f for f in os.listdir(data_resource_path) if f.endswith('.json')] if not json_files: raise ValueError(f"在 {data_resource_path} 目录下没有找到json文件") args.input = os.path.join(data_resource_path, random.choice(json_files)) print(f"随机选择输入文件: {args.input}") args.output = "tmp.json" print(f"使用默认输出文件: {args.output}") return args def main(): args = parse_args() modules_to_run = None if args.modules: modules_to_run = [m.strip() for m in args.modules] ok = run_pipeline( input_path=args.input, output_path=args.output, temp_dir=args.temp_dir, modules_to_run=modules_to_run, threads=args.threads, chart_name=args.chart_name, chart_only=args.chart_only, output_png=args.output_png, ) raise SystemExit(0 if ok else 1) if __name__ == "__main__": main()