179 lines
6.7 KiB
Python
179 lines
6.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
|
|
from .io_gltf import GltfDocument
|
|
from .accessors import access_data, update_accessor_binary_data, recompute_accessor_min_max
|
|
|
|
|
|
def quaternion_matrix(quaternion: np.ndarray) -> np.ndarray:
|
|
x, y, z, w = quaternion
|
|
x2 = x * x
|
|
y2 = y * y
|
|
z2 = z * z
|
|
xy = x * y
|
|
xz = x * z
|
|
yz = y * z
|
|
wx = w * x
|
|
wy = w * y
|
|
wz = w * z
|
|
return np.array(
|
|
[
|
|
[1.0 - 2.0 * (y2 + z2), 2.0 * (xy - wz), 2.0 * (xz + wy), 0.0],
|
|
[2.0 * (xy + wz), 1.0 - 2.0 * (x2 + z2), 2.0 * (yz - wx), 0.0],
|
|
[2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (x2 + y2), 0.0],
|
|
[0.0, 0.0, 0.0, 1.0],
|
|
]
|
|
)
|
|
|
|
|
|
def compose_matrix(scale: np.ndarray, rotation: np.ndarray, translation: np.ndarray) -> np.ndarray:
|
|
S = np.diag([scale[0], scale[1], scale[2], 1.0])
|
|
R = quaternion_matrix(rotation)
|
|
T = np.eye(4)
|
|
T[:3, 3] = translation
|
|
return T @ (R @ S)
|
|
|
|
|
|
def direction_to_quaternion(direction: np.ndarray) -> np.ndarray:
|
|
if not np.issubdtype(direction.dtype, np.floating):
|
|
direction = direction.astype(np.float32)
|
|
norm = np.linalg.norm(direction)
|
|
if norm < 1e-10:
|
|
return np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
direction = direction / norm
|
|
forward = np.array([0.0, 1.0, 0.0], dtype=np.float32)
|
|
up = np.array([0.0, 0.0, 1.0], dtype=np.float32)
|
|
dot = float(np.clip(np.dot(forward, direction), -1.0, 1.0))
|
|
if dot > 0.9999:
|
|
return np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
|
if dot < -0.9999:
|
|
return np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
|
|
axis = np.cross(forward, direction)
|
|
axis_norm = np.linalg.norm(axis)
|
|
if axis_norm < 1e-10:
|
|
if abs(direction[2]) < 0.9:
|
|
axis = np.cross(up, direction)
|
|
else:
|
|
axis = np.cross(np.array([1.0, 0.0, 0.0], dtype=np.float32), direction)
|
|
axis = axis / np.linalg.norm(axis)
|
|
else:
|
|
axis = axis / axis_norm
|
|
angle = np.arccos(dot)
|
|
half = angle * 0.5
|
|
s = np.sin(half)
|
|
c = np.cos(half)
|
|
quat = np.array([axis[0] * s, axis[1] * s, axis[2] * s, c], dtype=np.float32)
|
|
return quat / np.linalg.norm(quat)
|
|
|
|
|
|
def parse_node_transform(node_json: Dict[str, Any]) -> np.ndarray:
|
|
if "matrix" in node_json:
|
|
m = np.array(node_json["matrix"], dtype=np.float32).reshape(4, 4)
|
|
return m.T
|
|
t = np.array(node_json.get("translation", [0.0, 0.0, 0.0]), dtype=np.float32)
|
|
r = np.array(node_json.get("rotation", [0.0, 0.0, 0.0, 1.0]), dtype=np.float32)
|
|
s = np.array(node_json.get("scale", [1.0, 1.0, 1.0]), dtype=np.float32)
|
|
return compose_matrix(s, r, t)
|
|
|
|
|
|
def compute_global_transform(doc: GltfDocument, node_index: int, parent_transform: Optional[np.ndarray] = None) -> np.ndarray:
|
|
node_json = doc.nodes()[node_index]
|
|
local = parse_node_transform(node_json)
|
|
return parent_transform @ local if parent_transform is not None else local
|
|
|
|
|
|
def compute_mesh_global_transforms(doc: GltfDocument) -> Dict[int, np.ndarray]:
|
|
mesh_to_transform: Dict[int, np.ndarray] = {}
|
|
scene_index = doc.json().get("scene")
|
|
if scene_index is None:
|
|
return mesh_to_transform
|
|
root_nodes = doc.scenes()[scene_index].get("nodes", [])
|
|
|
|
def traverse(node_idx: int, parent: Optional[np.ndarray] = None) -> None:
|
|
global_m = compute_global_transform(doc, node_idx, parent)
|
|
node = doc.nodes()[node_idx]
|
|
if "mesh" in node:
|
|
mesh_to_transform[int(node["mesh"])] = global_m
|
|
for child in node.get("children", []) or []:
|
|
traverse(int(child), global_m)
|
|
|
|
for root in root_nodes:
|
|
traverse(int(root))
|
|
return mesh_to_transform
|
|
|
|
|
|
def transform_positions_inplace(doc: GltfDocument, mesh_idx: int, matrix: np.ndarray) -> None:
|
|
mesh = doc.meshes()[mesh_idx]
|
|
for primitive in mesh.get("primitives", []):
|
|
attrs = primitive.get("attributes", {})
|
|
if "POSITION" not in attrs:
|
|
continue
|
|
acc_idx = int(attrs["POSITION"])
|
|
data = access_data(doc, acc_idx)
|
|
pos_h = np.ones((data.shape[0], 4), dtype=np.float32)
|
|
pos_h[:, :3] = data.astype(np.float32)
|
|
transformed = (pos_h @ matrix.T)[:, :3].astype(np.float32)
|
|
update_accessor_binary_data(doc, acc_idx, transformed)
|
|
recompute_accessor_min_max(doc, acc_idx)
|
|
|
|
|
|
def transform_local_vertices_to_world(pos_data: np.ndarray, node: Dict[str, Any]) -> np.ndarray:
|
|
if "matrix" in node:
|
|
M = np.array(node["matrix"], dtype=np.float32).reshape(4, 4).T
|
|
else:
|
|
t = np.array(node.get("translation", [0.0, 0.0, 0.0]), dtype=np.float32)
|
|
r = np.array(node.get("rotation", [0.0, 0.0, 0.0, 1.0]), dtype=np.float32)
|
|
s = np.array(node.get("scale", [1.0, 1.0, 1.0]), dtype=np.float32)
|
|
M = compose_matrix(s, r, t)
|
|
pos_h = np.concatenate([pos_data, np.ones((pos_data.shape[0], 1), dtype=np.float32)], axis=1)
|
|
out = (M @ pos_h.T).T[:, :3]
|
|
return out
|
|
|
|
|
|
def set_node_trs_identity(doc: GltfDocument, node_idx: int) -> None:
|
|
node = doc.nodes()[node_idx]
|
|
for key in ["matrix", "translation", "rotation", "scale"]:
|
|
if key in node:
|
|
node.pop(key)
|
|
|
|
|
|
def preprocess_transforms_inplace(doc: GltfDocument) -> None:
|
|
"""Apply node transforms to vertex positions and set TRS for all nodes.
|
|
|
|
Idea: for each mesh, we calculate the global matrix, transform the POSITION of all
|
|
primitives of this mesh and remove the transform for the nodes referring to this mesh.
|
|
"""
|
|
mesh_transforms = compute_mesh_global_transforms(doc)
|
|
# quick exit if all identity
|
|
any_non_identity = any(
|
|
not np.allclose(M, np.eye(4), atol=1e-6) for M in mesh_transforms.values()
|
|
)
|
|
if not any_non_identity:
|
|
return
|
|
|
|
for mesh_idx, M in mesh_transforms.items():
|
|
if np.allclose(M, np.eye(4), atol=1e-6):
|
|
continue
|
|
# transform all positions in the mesh
|
|
mesh = doc.meshes()[mesh_idx]
|
|
for primitive in mesh.get("primitives", []):
|
|
attrs = primitive.get("attributes", {})
|
|
if "POSITION" not in attrs:
|
|
continue
|
|
acc_idx = int(attrs["POSITION"])
|
|
data = access_data(doc, acc_idx)
|
|
pos_h = np.ones((data.shape[0], 4), dtype=np.float32)
|
|
pos_h[:, :3] = data.astype(np.float32)
|
|
transformed = (pos_h @ M.T)[:, :3].astype(np.float32)
|
|
update_accessor_binary_data(doc, acc_idx, transformed)
|
|
recompute_accessor_min_max(doc, acc_idx)
|
|
|
|
# reset TRS for all nodes referring to this mesh
|
|
for node_idx, node in enumerate(doc.nodes()):
|
|
if node.get("mesh") == mesh_idx:
|
|
set_node_trs_identity(doc, node_idx)
|
|
|