blender-vrm/editor/body_mesh/utils.py

# SPDX-License-Identifier: MIT OR GPL-3.0-or-later
import json
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from bmesh.types import BMesh, BMLayerItem
    from bpy.types import Context, Mesh, Object


# UV坐标精度（小数位数）
UV_PRECISION = 6


# 材质到网格的映射
MATERIAL_MESH_MAPPING: dict[str, str] = {
    "N00_000_00_Body_00_SKIN (Instance)": "Body",
    "N00_000_00_Face_00_SKIN (Instance)": "Face",
    "N00_000_00_FaceMouth_00_FACE (Instance)": "Face",
    "N00_000_00_FaceEyeline_00_FACE (Instance)": "Face",
    "N00_000_00_FaceEyelash_00_FACE (Instance)": "Face",
    "N00_000_00_FaceBrow_00_FACE (Instance)": "Face",
    "N00_000_00_hH_00_FACE (Instance)": "Face",
    "N00_000_00_EyeIris_00_EYE (Instance)": "Face",
    "N00_000_00_EyeWhite_00_EYE (Instance)": "Face",
    "N00_000_00_EyeHighlight_00_EYE (Instance)": "Face",
    "N00_000_00_HairBack_00_HAIR (Instance)": "Body",
}

# 使用Blender自带tris_convert_to_quads的材质（网格规律，不需要UV匹配）
AUTO_CONVERT_MATERIALS: set[str] = {
    "N00_000_00_FaceBrow_00_FACE (Instance)",
}

# 处理顺序（先Body后Face）
PROCESSING_ORDER: list[str] = [
    "N00_000_00_Body_00_SKIN (Instance)",
    "N00_000_00_HairBack_00_HAIR (Instance)",
    "N00_000_00_Face_00_SKIN (Instance)",
    "N00_000_00_FaceMouth_00_FACE (Instance)",
    "N00_000_00_FaceEyeline_00_FACE (Instance)",
    "N00_000_00_FaceEyelash_00_FACE (Instance)",
    "N00_000_00_FaceBrow_00_FACE (Instance)",
    "N00_000_00_hH_00_FACE (Instance)",
    "N00_000_00_EyeIris_00_EYE (Instance)",
    "N00_000_00_EyeWhite_00_EYE (Instance)",
    "N00_000_00_EyeHighlight_00_EYE (Instance)",
]


@dataclass
class MaterialInfo:
    """材质信息数据类"""
    name: str
    index: int
    face_count: int
    is_skin: bool


def find_body_mesh(armature: "Object") -> Optional["Object"]:
    """
    在Armature的子对象中查找名为"Body"的网格对象

    Args:
        armature: Armature对象

    Returns:
        Body网格对象，如果不存在则返回None
    """
    if armature is None:
        return None

    for child in armature.children:
        if child.type == "MESH" and child.name == "Body":
            return child

    return None


def find_face_mesh(armature: "Object") -> Optional["Object"]:
    """
    在Armature的子对象中查找名为"Face"的网格对象

    Args:
        armature: Armature对象

    Returns:
        Face网格对象，如果不存在则返回None
    """
    if armature is None:
        return None

    for child in armature.children:
        if child.type == "MESH" and child.name == "Face":
            return child

    return None


def get_target_mesh_name(material_name: str) -> Optional[str]:
    """
    根据材质名称获取目标网格名称

    Args:
        material_name: 材质名称

    Returns:
        目标网格名称（"Body"或"Face"），如果材质不在映射中返回None
    """
    return MATERIAL_MESH_MAPPING.get(material_name)


def get_material_index(mesh_obj: "Object", material_name: str) -> Optional[int]:
    """
    获取材质在网格中的索引

    Args:
        mesh_obj: 网格对象
        material_name: 材质名称

    Returns:
        材质索引，如果不存在返回None
    """
    if mesh_obj is None or mesh_obj.type != "MESH":
        return None

    for i, slot in enumerate(mesh_obj.material_slots):
        if slot.material and slot.material.name == material_name:
            return i

    return None


# 保留材质关键字常量（不分离的材质：只有SKIN和HAIR）
RETAINED_MATERIAL_KEYWORDS = ["SKIN", "HAIR"]


def is_skin_material(material_name: str) -> bool:
    """
    判断材质是否为皮肤材质

    Args:
        material_name: 材质名称

    Returns:
        如果材质名包含"SKIN"则返回True（不区分大小写）
    """
    return "SKIN" in material_name.upper()


def is_retained_material(material_name: str) -> bool:
    """
    判断材质是否为保留材质（不分离）

    只有SKIN和HAIR材质保留在原Body网格中。
    CLOTH和其他材质都需要分离。

    Args:
        material_name: 材质名称

    Returns:
        如果材质名包含"SKIN"或"HAIR"则返回True（不区分大小写）
    """
    upper_name = material_name.upper()
    return any(keyword in upper_name for keyword in RETAINED_MATERIAL_KEYWORDS)


def is_separable_material(material_name: str) -> bool:
    """
    判断材质是否为可分离材质

    不包含"SKIN"、"HAIR"的材质都是可分离材质，
    包括CLOTH和其他所有材质。

    Args:
        material_name: 材质名称

    Returns:
        如果材质名不包含"SKIN"、"HAIR"则返回True（包括CLOTH和其他材质）
    """
    return not is_retained_material(material_name)


def get_material_face_count(mesh: "Mesh", material_index: int) -> int:
    """
    获取指定材质的面数量

    Args:
        mesh: 网格数据
        material_index: 材质索引

    Returns:
        使用该材质的面数量
    """
    count = 0
    for polygon in mesh.polygons:
        if polygon.material_index == material_index:
            count += 1
    return count


def generate_unique_name(base_name: str, existing_names: set[str]) -> str:
    """
    生成唯一的对象名称

    Args:
        base_name: 基础名称
        existing_names: 已存在的名称集合

    Returns:
        唯一的名称
    """
    if base_name not in existing_names:
        return base_name

    suffix = 1
    while True:
        candidate = f"{base_name}.{suffix:03d}"
        if candidate not in existing_names:
            return candidate
        suffix += 1


def normalize_uv_pair(
    uv1: tuple[float, float], uv2: tuple[float, float]
) -> tuple[tuple[float, float], tuple[float, float]]:
    """
    标准化UV坐标对（较小的在前）

    通过将较小的UV坐标放在前面，确保相同的边无论顶点顺序如何都能匹配。
    使用6位小数精度进行比较和存储。

    Args:
        uv1: 第一个UV坐标 (u, v)
        uv2: 第二个UV坐标 (u, v)

    Returns:
        标准化后的UV对 ((u1,v1), (u2,v2))，较小的坐标在前
    """
    # 四舍五入到指定精度
    uv1_rounded = (round(uv1[0], UV_PRECISION), round(uv1[1], UV_PRECISION))
    uv2_rounded = (round(uv2[0], UV_PRECISION), round(uv2[1], UV_PRECISION))

    # 比较并排序，较小的在前
    if uv1_rounded <= uv2_rounded:
        return (uv1_rounded, uv2_rounded)
    return (uv2_rounded, uv1_rounded)


def load_edge_uv_data(material_name: str) -> Optional[set[tuple]]:
    """
    从JSON文件加载边UV数据

    Args:
        material_name: 材质名称（用于定位JSON文件）

    Returns:
        UV边坐标集合，格式为 {((u1,v1), (u2,v2)), ...}
        如果文件不存在返回None
    """
    # 获取当前模块所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    json_path = os.path.join(current_dir, f"{material_name}.json")

    if not os.path.exists(json_path):
        return None

    try:
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        edge_uvs = data.get("edge_uvs", [])
        result: set[tuple] = set()

        for edge_uv in edge_uvs:
            if len(edge_uv) == 2 and len(edge_uv[0]) == 2 and len(edge_uv[1]) == 2:
                uv1 = (float(edge_uv[0][0]), float(edge_uv[0][1]))
                uv2 = (float(edge_uv[1][0]), float(edge_uv[1][1]))
                normalized = normalize_uv_pair(uv1, uv2)
                result.add(normalized)

        return result
    except (json.JSONDecodeError, KeyError, TypeError, ValueError):
        return None


def select_edges_by_uv(
    bm: "BMesh",
    uv_layer: "BMLayerItem",
    target_uvs: set[tuple],
) -> int:
    """
    根据UV坐标选中边

    遍历BMesh中的所有边，比较每条边的loop UV坐标与目标UV集合，
    选中匹配的边。使用6位小数精度进行匹配。

    Args:
        bm: BMesh对象
        uv_layer: UV层
        target_uvs: 目标UV坐标集合，格式为 {((u1,v1), (u2,v2)), ...}

    Returns:
        选中的边数量
    """
    selected_count = 0

    for edge in bm.edges:
        # 获取边的两个顶点关联的loop
        # 每条边可能有多个关联的面，我们需要检查所有关联的loop
        edge_selected = False

        for face in edge.link_faces:
            # 找到这条边在当前面中的两个loop
            edge_loops = []
            for loop in face.loops:
                if loop.vert in edge.verts:
                    edge_loops.append(loop)

            if len(edge_loops) == 2:
                # 获取两个loop的UV坐标
                uv1_data = edge_loops[0][uv_layer]
                uv2_data = edge_loops[1][uv_layer]

                uv1 = (uv1_data.uv[0], uv1_data.uv[1])
                uv2 = (uv2_data.uv[0], uv2_data.uv[1])

                # 标准化UV对并检查是否在目标集合中
                normalized = normalize_uv_pair(uv1, uv2)

                if normalized in target_uvs:
                    edge_selected = True
                    break

        if edge_selected:
            edge.select = True
            selected_count += 1

    return selected_count


def isolate_material_faces(
    context: "Context",
    mesh_obj: "Object",
    material_name: str,
) -> bool:
    """
    隔离显示指定材质的面（隐藏其他面）

    进入编辑模式，选择指定材质的所有面，执行隐藏未选中操作，
    然后切换到边选择模式。

    Args:
        context: Blender上下文
        mesh_obj: 网格对象
        material_name: 材质名称

    Returns:
        是否成功隔离

    Requirements: 3.1, 3.2, 3.3, 3.4
    """
    import bpy

    if mesh_obj is None or mesh_obj.type != "MESH":
        return False

    # 获取材质索引
    material_index = get_material_index(mesh_obj, material_name)
    if material_index is None:
        return False

    # 保存当前活动对象和选择状态
    original_active = context.view_layer.objects.active
    original_mode = context.object.mode if context.object else "OBJECT"

    try:
        # 确保在对象模式下开始
        if original_mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        # 取消所有选择，选择目标网格
        bpy.ops.object.select_all(action="DESELECT")
        mesh_obj.select_set(True)
        context.view_layer.objects.active = mesh_obj

        # 进入编辑模式 (Requirements 3.1)
        bpy.ops.object.mode_set(mode="EDIT")

        # 切换到面选择模式
        bpy.ops.mesh.select_mode(type="FACE")

        # 取消所有选择
        bpy.ops.mesh.select_all(action="DESELECT")

        # 设置活动材质索引并选择该材质的所有面 (Requirements 3.2)
        mesh_obj.active_material_index = material_index
        bpy.ops.object.material_slot_select()

        # 执行隐藏未选中 (Shift+H equivalent) (Requirements 3.3)
        bpy.ops.mesh.hide(unselected=True)

        # 切换到边选择模式 (Requirements 3.4)
        bpy.ops.mesh.select_mode(type="EDGE")

        return True

    except Exception:
        # 发生错误时尝试恢复状态
        try:
            if context.object and context.object.mode != "OBJECT":
                bpy.ops.object.mode_set(mode="OBJECT")
            if original_active:
                context.view_layer.objects.active = original_active
        except Exception:
            pass
        return False


def select_material_faces(
    mesh_obj: "Object",
    material_index: int,
) -> set[int]:
    """
    获取指定材质索引的所有面的索引集合

    这是一个纯函数，用于测试材质面选择的完整性。
    不修改任何Blender状态。

    Args:
        mesh_obj: 网格对象
        material_index: 材质索引

    Returns:
        使用该材质的面索引集合
    """
    if mesh_obj is None or mesh_obj.type != "MESH":
        return set()

    mesh = mesh_obj.data
    result: set[int] = set()

    for i, polygon in enumerate(mesh.polygons):
        if polygon.material_index == material_index:
            result.add(i)

    return result