starmen1265
042bebe9d7
- 默认模型改为 glm-5.1，API地址改为 coding 套餐端点
- 默认批次大小 20→5，增加 max_tokens=4096
- 优化 prompt 模板，明确要求返回 {"results": [...]} 格式
- 修复全批失败时 ai_inference_metadata 的 KeyError
- 输出目录从 <project>/output/ 统一到 <tool>/outputs/
- 新增 README.md（Pipeline流程、环境变量、CLI参数说明）
370 lines
12 KiB
Python
370 lines
12 KiB
Python
"""GLM API 批量推理脚本 — 资产语义增强。
|
||
|
||
读取 asset_database.json，按类别分批调用 GLM API（OpenAI 兼容接口），
|
||
推断每个资产的朝向、功能、交互属性等语义信息。
|
||
|
||
此脚本不依赖 UE5，可离线运行。
|
||
|
||
环境变量:
|
||
GLM_API_KEY — API Key（必需）
|
||
GLM_API_BASE — API Base URL（默认 https://open.bigmodel.cn/api/paas/v4）
|
||
GLM_MODEL — 模型名（默认 glm-4-flash）
|
||
|
||
用法:
|
||
python ai_inference.py [--input FILE] [--output FILE] [--batch-size N] [--dry-run]
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import json
|
||
import argparse
|
||
import time
|
||
from pathlib import Path
|
||
|
||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||
if str(_SCRIPT_DIR) not in sys.path:
|
||
sys.path.insert(0, str(_SCRIPT_DIR))
|
||
|
||
from _common import read_json, write_json, iso_now
|
||
|
||
TAG = "[AIInference]"
|
||
|
||
# ---- 配置 ----
|
||
DEFAULT_API_BASE = "https://open.bigmodel.cn/api/coding/paas/v4"
|
||
DEFAULT_MODEL = "glm-5.1"
|
||
DEFAULT_BATCH_SIZE = 5
|
||
|
||
# ---- Prompt 模板 ----
|
||
SYSTEM_PROMPT = """你是一个 UE5 游戏开发助手，专注于为 3D 模型资产提取语义信息。
|
||
|
||
你会收到一批 StaticMesh 资产的几何数据（名称、尺寸、锚点、材质等），
|
||
需要推断每个资产的以下语义字段：
|
||
|
||
1. facing_direction: 主朝向方向
|
||
- 格式: "+X", "-X", "+Y", "-Y", null
|
||
- 判断依据: 根据资产类型和包围盒尺寸推断"正面"方向
|
||
- 对称物体（桌子、球、圆柱等）设为 null
|
||
- 电视/显示器: 屏幕面方向; 门: 开门方向; 椅子: 坐面方向
|
||
|
||
2. surface_height_cm: 可用水平表面高度 (cm)
|
||
- 如桌面高度、台面高度、床面高度
|
||
- 对于无水平表面的物体（墙体、栏杆）设为 null
|
||
- 使用锚点类型判断: CENTER_BOTTOM 时直接取 max.z
|
||
|
||
3. function_description: 一句话功能描述 (中文, ≤50字)
|
||
|
||
4. placement_notes: 放置注意事项 (中文, ≤80字)
|
||
- 锚点位置提醒、旋转建议、缩放注意等
|
||
|
||
5. interaction_hints: 交互提示标签数组
|
||
- 可选值: can_sit, can_place_items, can_open, can_interact_screen,
|
||
climbable, electronics, lighting, container, decoration,
|
||
pass_through, driveable, weapon, vegetation
|
||
- 可多选
|
||
|
||
6. plausible_substitutes: BSP 基础体降级方案 (中文, 或 null)
|
||
- 当原始资产不可用时，用什么 Cube/Cylinder + 缩放 + 材质替代
|
||
|
||
你必须返回严格的 JSON 数组，每个元素对应输入中的一个资产，顺序一致。
|
||
不要输出任何 JSON 之外的内容。"""
|
||
|
||
BATCH_PROMPT_TEMPLATE = """以下是 {category} 类别的 {count} 个 StaticMesh 资产数据：
|
||
|
||
{assets_json}
|
||
|
||
请为每个资产推断语义信息，必须返回如下格式的 JSON 对象：
|
||
{{"results": [{{"facing_direction": ..., "surface_height_cm": ..., "function_description": ..., "placement_notes": ..., "interaction_hints": [...], "plausible_substitutes": ...}}, ...]}}
|
||
重要：results 数组必须恰好包含 {count} 个元素，与输入资产一一对应。"""
|
||
|
||
|
||
def _get_api_config():
|
||
"""从环境变量读取 API 配置。"""
|
||
api_key = os.environ.get("GLM_API_KEY", "")
|
||
api_base = os.environ.get("GLM_API_BASE", DEFAULT_API_BASE)
|
||
model = os.environ.get("GLM_MODEL", DEFAULT_MODEL)
|
||
|
||
if not api_key:
|
||
print(f"{TAG} 错误: 未设置 GLM_API_KEY 环境变量")
|
||
print(f"{TAG} 设置方法: set GLM_API_KEY=your_key_here")
|
||
return None
|
||
|
||
return {
|
||
"api_key": api_key,
|
||
"api_base": api_base,
|
||
"model": model,
|
||
}
|
||
|
||
|
||
def _classify_batch_category(tags: dict) -> str:
|
||
"""根据 tags 推断批次类别名。"""
|
||
style = tags.get("style", [])
|
||
if not style:
|
||
return "未分类"
|
||
# 取第一个非空 style tag 作为类别
|
||
return style[0] if style[0] != "unknown" else "未分类"
|
||
|
||
|
||
def _prepare_batches(assets: dict, batch_size: int) -> list:
|
||
"""将资产按类别分批。
|
||
|
||
Returns:
|
||
list of (category_name, batch_index, [asset_keys])
|
||
"""
|
||
# 按类别分组
|
||
by_category = {}
|
||
for name, data in assets.items():
|
||
cat = _classify_batch_category(data.get("tags", {}))
|
||
by_category.setdefault(cat, []).append(name)
|
||
|
||
# 每个类别内按 batch_size 切分
|
||
batches = []
|
||
for cat, keys in sorted(by_category.items()):
|
||
for i in range(0, len(keys), batch_size):
|
||
batch_keys = keys[i:i + batch_size]
|
||
batches.append((cat, i // batch_size + 1, batch_keys))
|
||
|
||
return batches
|
||
|
||
|
||
def _format_asset_for_prompt(name: str, data: dict) -> dict:
|
||
"""提取 prompt 中需要的精简字段。"""
|
||
bb = data.get("bounding_box", {})
|
||
pivot = data.get("pivot", {})
|
||
tags = data.get("tags", {})
|
||
|
||
return {
|
||
"name": name,
|
||
"size_cm": bb.get("size"),
|
||
"pivot_type": pivot.get("type"),
|
||
"pivot_offset_xy": pivot.get("offset_xy"),
|
||
"material_count": len(data.get("material_slots", [])),
|
||
"socket_count": len(data.get("sockets", [])),
|
||
"tags_style": tags.get("style", []),
|
||
}
|
||
|
||
|
||
def _call_glm_api(config: dict, category: str, assets_json: str) -> list:
|
||
"""调用 GLM API 进行推理。
|
||
|
||
Returns:
|
||
推理结果列表，失败时返回空列表
|
||
"""
|
||
try:
|
||
from openai import OpenAI
|
||
except ImportError:
|
||
print(f"{TAG} 错误: 需要安装 openai 包")
|
||
print(f"{TAG} pip install openai")
|
||
return []
|
||
|
||
client = OpenAI(
|
||
api_key=config["api_key"],
|
||
base_url=config["api_base"],
|
||
)
|
||
|
||
user_prompt = BATCH_PROMPT_TEMPLATE.format(
|
||
category=category,
|
||
count=assets_json.count('"name"'),
|
||
assets_json=assets_json,
|
||
)
|
||
|
||
try:
|
||
response = client.chat.completions.create(
|
||
model=config["model"],
|
||
messages=[
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": user_prompt},
|
||
],
|
||
temperature=0.1,
|
||
max_tokens=4096,
|
||
response_format={"type": "json_object"},
|
||
)
|
||
except Exception as e:
|
||
print(f"{TAG} API 调用失败: {e}")
|
||
return []
|
||
|
||
content = response.choices[0].message.content.strip()
|
||
if not content:
|
||
print(f"{TAG} API 返回空内容")
|
||
return []
|
||
|
||
# 解析 JSON
|
||
try:
|
||
result = json.loads(content)
|
||
# 如果返回的是 {"results": [...]} 形式，提取数组
|
||
if isinstance(result, dict):
|
||
for key in ("results", "data", "assets", "items"):
|
||
if key in result and isinstance(result[key], list):
|
||
return result[key]
|
||
# 如果找不到数组，可能是单条结果被包在 dict 里
|
||
print(f"{TAG} 警告: 返回 JSON 不是数组，尝试整体作为单条: {list(result.keys())}")
|
||
return [result]
|
||
if isinstance(result, list):
|
||
return result
|
||
except json.JSONDecodeError as e:
|
||
print(f"{TAG} JSON 解析失败: {e}")
|
||
print(f"{TAG} 原始响应前 200 字: {content[:200]}")
|
||
return []
|
||
|
||
return []
|
||
|
||
|
||
def run_inference(
|
||
input_file: Path,
|
||
output_file: Path,
|
||
batch_size: int = DEFAULT_BATCH_SIZE,
|
||
dry_run: bool = False,
|
||
) -> dict:
|
||
"""执行批量 AI 推理。
|
||
|
||
Args:
|
||
input_file: asset_database.json 路径
|
||
output_file: 输出路径（通常与 input 相同，就地更新）
|
||
batch_size: 每批资产数量
|
||
dry_run: 只打印批次信息，不调用 API
|
||
|
||
Returns:
|
||
统计信息 dict
|
||
"""
|
||
# 读取数据库
|
||
if not input_file.exists():
|
||
print(f"{TAG} 错误: {input_file} 不存在，请先运行 build_database")
|
||
return {"status": "error", "reason": "no input file"}
|
||
|
||
db = read_json(input_file)
|
||
assets = db.get("assets", {})
|
||
if not assets:
|
||
print(f"{TAG} 错误: 数据库中无资产")
|
||
return {"status": "error", "reason": "empty assets"}
|
||
|
||
# 过滤: 跳过已推理的
|
||
pending = {}
|
||
for name, data in assets.items():
|
||
if data.get("ai_inference") is not None:
|
||
continue
|
||
pending[name] = data
|
||
|
||
print(f"{TAG} 数据库共 {len(assets)} 个资产，待推理 {len(pending)} 个")
|
||
|
||
if not pending:
|
||
print(f"{TAG} 全部已推理，无需处理")
|
||
return {"status": "ok", "total": len(assets), "inferred": 0, "skipped": len(assets)}
|
||
|
||
# 分批
|
||
batches = _prepare_batches(pending, batch_size)
|
||
print(f"{TAG} 分为 {len(batches)} 批")
|
||
|
||
if dry_run:
|
||
for cat, idx, keys in batches:
|
||
print(f"{TAG} 批次 [{cat}] #{idx}: {len(keys)} 个资产")
|
||
for k in keys[:3]:
|
||
print(f"{TAG} - {k}")
|
||
if len(keys) > 3:
|
||
print(f"{TAG} ... 还有 {len(keys) - 3} 个")
|
||
return {"status": "ok", "dry_run": True, "batches": len(batches)}
|
||
|
||
# 获取 API 配置
|
||
config = _get_api_config()
|
||
if not config:
|
||
return {"status": "error", "reason": "no api key"}
|
||
|
||
print(f"{TAG} API: {config['api_base']} 模型: {config['model']}")
|
||
|
||
# 逐批推理
|
||
total_inferred = 0
|
||
total_failed = 0
|
||
|
||
for batch_idx, (cat, batch_num, keys) in enumerate(batches):
|
||
print(f"{TAG} 批次 [{batch_idx + 1}/{len(batches)}] [{cat}] #{batch_num}: {len(keys)} 个资产...")
|
||
|
||
# 构造 prompt 数据
|
||
prompt_assets = []
|
||
for k in keys:
|
||
prompt_assets.append(_format_asset_for_prompt(k, pending[k]))
|
||
|
||
assets_json = json.dumps(prompt_assets, ensure_ascii=False, indent=2)
|
||
|
||
# 调用 API
|
||
results = _call_glm_api(config, cat, assets_json)
|
||
|
||
if not results:
|
||
print(f"{TAG} ⚠️ 推理失败，跳过此批")
|
||
total_failed += len(keys)
|
||
continue
|
||
|
||
# 合并结果
|
||
matched = 0
|
||
for i, k in enumerate(keys):
|
||
if i < len(results):
|
||
assets[k]["ai_inference"] = results[i]
|
||
matched += 1
|
||
total_inferred += 1
|
||
else:
|
||
print(f"{TAG} ⚠️ {k}: API 返回结果不足")
|
||
|
||
print(f"{TAG} ✅ 匹配 {matched}/{len(keys)}")
|
||
|
||
# 每批都保存一次（断点续跑）
|
||
db["assets"] = assets
|
||
# 更新 metadata
|
||
if "ai_inference_metadata" not in db["metadata"]:
|
||
db["metadata"]["ai_inference_metadata"] = {
|
||
"model": config["model"],
|
||
"api_base": config["api_base"],
|
||
"started_at": iso_now(),
|
||
}
|
||
db["metadata"]["ai_inference_metadata"]["last_batch"] = batch_idx + 1
|
||
db["metadata"]["ai_inference_metadata"]["total_batches"] = len(batches)
|
||
db["metadata"]["ai_inference_metadata"]["inferred"] = total_inferred
|
||
db["metadata"]["ai_inference_metadata"]["failed"] = total_failed
|
||
write_json(db, output_file)
|
||
|
||
# 限速: 避免触发 API 限流
|
||
if batch_idx < len(batches) - 1:
|
||
time.sleep(1)
|
||
|
||
if "ai_inference_metadata" not in db["metadata"]:
|
||
db["metadata"]["ai_inference_metadata"] = {}
|
||
db["metadata"]["ai_inference_metadata"]["completed_at"] = iso_now()
|
||
write_json(db, output_file)
|
||
|
||
print(f"{TAG} 推理完成: {total_inferred} 成功, {total_failed} 失败")
|
||
return {
|
||
"status": "ok",
|
||
"total": len(assets),
|
||
"inferred": total_inferred,
|
||
"failed": total_failed,
|
||
}
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="GLM API 批量推理 — 资产语义增强")
|
||
parser.add_argument(
|
||
"--input", type=str,
|
||
default=str(_SCRIPT_DIR / "outputs" / "asset_database.json"),
|
||
help="输入 asset_database.json 路径",
|
||
)
|
||
parser.add_argument(
|
||
"--output", type=str, default=None,
|
||
help="输出路径 (默认与 input 相同，就地更新)",
|
||
)
|
||
parser.add_argument(
|
||
"--batch-size", type=int, default=DEFAULT_BATCH_SIZE,
|
||
help=f"每批资产数量 (默认 {DEFAULT_BATCH_SIZE})",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run", action="store_true",
|
||
help="只打印批次信息，不调用 API",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
input_file = Path(args.input)
|
||
output_file = Path(args.output) if args.output else input_file
|
||
|
||
result = run_inference(input_file, output_file, args.batch_size, args.dry_run)
|
||
if result.get("status") != "ok" and not result.get("dry_run"):
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|