ChildHood/tools/asset-catalog/discover_assets.py

"""UE5 资产发现脚本。

扫描指定 Content 路径下所有资产，记录名称、路径、UE5 类型。
不做任何推断，仅记录 UE5 原始数据。

输出: output/index.json

用法:
  交互模式: exec(open(r"...\discover_assets.py").read())
  命令行:   通过 run_all.py 调用
"""

import sys
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 (
    HAS_UNREAL, get_scan_root, get_output_dir, write_json, read_json,
    wait_for_editor, progress_log, iso_now, safe_subfolder_name,
)

TAG = "[Discover]"

# 已知类型分类
MESH_TYPES = {"StaticMesh"}
MATERIAL_TYPES = {"MaterialInstanceConstant", "Material"}
# 其余全部归入 others


def _get_subfolder(asset_path: str, scan_root: str) -> str:
    """从资产完整路径中提取相对于 scan_root/Meshes|Materials|... 的子目录。

    例:
      path = "/Game/BlockingStarterPack/Meshes/Architecture/Walls/SM_Wall_400x300.SM_Wall_400x300"
      scan_root = "/Game/BlockingStarterPack"
      → "Architecture/Walls"
    """
    # 去掉末尾的 .XXX
    base = asset_path.rsplit(".", 1)[0]
    # 找到 scan_root 之后的路径
    if base.startswith(scan_root + "/"):
        relative = base[len(scan_root) + 1:]
    else:
        relative = base

    # 去掉第一级目录和最后一级（资产名）
    # Meshes/Architecture/Walls/SM_Wall_400x300 → Architecture/Walls
    parts = relative.split("/")
    if len(parts) > 2:
        return "/".join(parts[1:-1])  # 去掉首段（Meshes 等）和末段（资产名）
    return ""


def _get_asset_name(asset_path: str) -> str:
    """从 UE 资产路径提取资产名。

    例: "/Game/.../SM_Wall_400x300.SM_Wall_400x300" → "SM_Wall_400x300"
    """
    # 取最后一个 / 后、第一个 . 前
    last_part = asset_path.rsplit("/", 1)[-1]
    return last_part.rsplit(".", 1)[0]


def _load_and_classify(asset_path: str) -> dict:
    """加载资产并记录其 UE5 类型名。

    Returns:
        {"name": str, "ue_path": str, "ue_class_type": str}
        加载失败时 ue_class_type 为 "LOAD_FAILED"
    """
    name = _get_asset_name(asset_path)
    try:
        import unreal
        obj = unreal.load_asset(asset_path)
        if obj is None:
            return {"name": name, "ue_path": asset_path, "ue_class_type": "LOAD_FAILED"}
        class_name = type(obj).__name__
        return {"name": name, "ue_path": asset_path, "ue_class_type": class_name}
    except Exception as e:
        return {"name": name, "ue_path": asset_path, "ue_class_type": f"ERROR:{e}"}


def run() -> dict:
    """执行资产发现。

    Returns:
        统计信息 dict
    """
    if not HAS_UNREAL:
        print(f"{TAG} 错误: unreal 模块不可用")
        return {"status": "error", "reason": "no unreal"}

    import unreal

    scan_root = get_scan_root()
    output_dir = get_output_dir()

    print(f"{TAG} 扫描根目录: {scan_root}")
    print(f"{TAG} 输出目录: {output_dir}")

    # ---- 枚举所有资产 ----
    print(f"{TAG} 开始枚举资产...")
    all_assets = []
    try:
        asset_list = unreal.EditorAssetLibrary.list_assets(scan_root, recursive=True, include_folder=False)
        all_assets = list(asset_list)
    except Exception as e:
        print(f"{TAG} 错误: list_assets 失败: {e}")
        return {"status": "error", "reason": str(e)}

    total = len(all_assets)
    print(f"{TAG} 发现 {total} 个资产")

    if total == 0:
        print(f"{TAG} 未发现任何资产，退出")
        return {"status": "ok", "total": 0}

    # ---- 分类 ----
    meshes = {}      # subfolder -> [asset_info, ...]
    materials = {}   # subfolder -> [asset_info, ...]
    others = {}      # subfolder -> [asset_info, ...]
    by_type = {}

    for i, asset_path in enumerate(all_assets):
        progress_log(i, total, "资产分类")

        info = _load_and_classify(asset_path)
        class_type = info["ue_class_type"]
        subfolder = _get_subfolder(asset_path, scan_root)

        # 统计
        by_type[class_type] = by_type.get(class_type, 0) + 1

        # 分类
        if class_type in MESH_TYPES:
            meshes.setdefault(subfolder, []).append(info)
        elif class_type in MATERIAL_TYPES:
            materials.setdefault(subfolder, []).append(info)
        else:
            others.setdefault(subfolder, []).append(info)

    # ---- 蓝图归类 ----
    # 从 others 中提取蓝图相关（Blueprint 类型或目录含 Blueprint 的）
    blueprints = {}
    # UE5 的蓝图资产 class type 可能是各种名称
    # 我们把 ue_path 中包含 /Blueprints/ 的单独归类
    remaining_others = {}
    for subfolder, assets in others.items():
        bp_assets = []
        other_assets = []
        for a in assets:
            if "/Blueprints/" in a["ue_path"] or a["ue_class_type"] in {
                "Blueprint", "BlueprintGeneratedClass", "WidgetBlueprint",
                "AnimBlueprint", "BlueprintType",
            }:
                bp_assets.append(a)
            else:
                other_assets.append(a)
        if bp_assets:
            blueprints[subfolder] = bp_assets
        if other_assets:
            remaining_others[subfolder] = other_assets

    # ---- 写入 index.json ----
    index = {
        "metadata": {
            "generated_at": iso_now(),
            "scan_root": scan_root,
            "total_assets": total,
            "by_type": by_type,
        },
        "meshes": meshes,
        "materials": materials,
        "blueprints": blueprints,
        "others": remaining_others,
    }

    output_path = output_dir / "index.json"
    write_json(index, output_path)
    print(f"{TAG} ✅ 已写入: {output_path}")

    # ---- 统计 ----
    mesh_count = sum(len(v) for v in meshes.values())
    mat_count = sum(len(v) for v in materials.values())
    bp_count = sum(len(v) for v in blueprints.values())
    other_count = sum(len(v) for v in remaining_others.values())

    stats = {
        "status": "ok",
        "total": total,
        "meshes": mesh_count,
        "materials": mat_count,
        "blueprints": bp_count,
        "others": other_count,
        "mesh_subfolders": list(meshes.keys()),
        "material_subfolders": list(materials.keys()),
        "blueprint_subfolders": list(blueprints.keys()),
    }

    print(f"{TAG} 统计:")
    print(f"{TAG}   StaticMesh: {mesh_count} ({len(meshes)} 子目录)")
    print(f"{TAG}   Materials:  {mat_count} ({len(materials)} 子目录)")
    print(f"{TAG}   Blueprints: {bp_count} ({len(blueprints)} 子目录)")
    print(f"{TAG}   Others:     {other_count}")

    return stats


# ---- 入口 ----
if __name__ == "__main__":
    if not wait_for_editor():
        sys.exit(1)
    result = run()
    if result.get("status") != "ok":
        sys.exit(1)