Move some prefab processing from JavaScript to Python script

This commit is contained in:
dominik-korsa 2025-07-19 02:13:17 +02:00
parent c7dec4abaf
commit 0d7e8dc18c
No known key found for this signature in database
GPG Key ID: 5A24D76EE0A30974

View File

@ -1,8 +1,19 @@
#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "pyyaml",
# "scipy",
# ]
# ///
import json
import os
import re
import sys
import traceback
from scipy.spatial.transform import Rotation
from pathlib import Path
import yaml
@ -98,38 +109,42 @@ def try_parse_track_transform(component, ancestor_level=0):
if "Transform" not in component:
return None
transform = component["Transform"]
parent = None
if "Transform" in transform["m_Father"]:
parent = try_parse_track_transform(transform["m_Father"], True)
else:
parent = {
"local_position": [0, 0, 0],
"local_rotation": Rotation.identity(),
"negate_radius": False,
}
if transform["m_LocalScale"]["x"] not in {-1, 1} or transform["m_LocalScale"]["y"] != 1 or transform["m_LocalScale"]["z"] != 1:
eprint("Unexpected scale in track transform")
if ancestor_level > 0:
if transform["m_LocalPosition"]["x"] != 0 or transform["m_LocalPosition"]["y"] != 0 or transform["m_LocalPosition"]["z"] != 0:
eprint("Unexpected position change in ancestor track transform")
if transform["m_LocalRotation"]["x"] != 0 or transform["m_LocalRotation"]["y"] != 0 or transform["m_LocalRotation"]["z"] != 0 or transform["m_LocalRotation"]["w"] != 1:
eprint("Unexpected rotation change in ancestor track transform")
local_position = [transform["m_LocalRotation"][key] for key in ["x", "y", "z"]]
local_rotation = Rotation.from_quat([transform["m_LocalRotation"][key] for key in ["x", "y", "z", "w"]])
negate_radius = parent["negate_radius"]
mirror_x = transform["m_LocalScale"]["x"] == -1
if mirror_x and ancestor_level != 1:
eprint("Expected mirror_x only in the direct parent of the track transform")
if transform["m_LocalScale"]["x"] == -1:
negate_radius = not negate_radius
rot_vet = local_rotation.as_rotvec()
# Flip the x-axis to apply the mirroring and then negate the whole vector to change the rotation direction
rot_vet[1] = -rot_vet[1]
rot_vet[2] = -rot_vet[2]
local_rotation = Rotation.from_rotvec(rot_vet)
parent_mirror_x = False
if parent is not None:
parent_mirror_x = parent["mirror_x"]
if parent is not None and parent["negate_radius"]:
negate_radius = not negate_radius
return {
"position": transform["m_LocalPosition"],
"rotation": transform["m_LocalRotation"],
"scale": transform["m_LocalScale"],
"mirror_x": mirror_x,
"parent_mirror_x": parent_mirror_x,
"parent": parent,
"local_position": parent["local_position"] + parent["local_rotation"].apply(local_position),
"local_rotation": parent["local_rotation"] * local_rotation,
"negate_radius": negate_radius,
}
def parse_track(track):
def parse_track(track, data_index):
track_id = track["__component_id"]
track_shape = None
track_transform = None
@ -146,21 +161,38 @@ def parse_track(track):
if maybe_track_transform is not None:
track_transform = maybe_track_transform
prev_id = None
next_id = None
if "__component_id" in track["MonoBehaviour"]["prevTrack"]:
prev_id = track["MonoBehaviour"]["prevTrack"]["__component_id"]
if "__component_id" in track["MonoBehaviour"]["nextTrack"]:
next_id = track["MonoBehaviour"]["nextTrack"]["__component_id"]
connections = []
def append_connection(end, field):
if "__component_id" in track["MonoBehaviour"][field]:
other_id = track["MonoBehaviour"][field]["__component_id"]
if other_id == track_id:
# Ignore connections to self, TD2 has them in Rkp switches
return
connections.append({
"type": "<INTERNAL>",
"end": end,
"internalId": other_id
})
else:
connections.append({
"type": "<EXTERNAL>",
"end": end,
})
append_connection("<START>", "prevTrack")
append_connection("<END>", "nextTrack")
if track_transform["negate_radius"]:
track_shape["radius"] = -track_shape["radius"]
return {
"id": track["__component_id"],
"prev_id": prev_id,
"next_id": next_id,
"shape": track_shape,
"mirror_x": track_transform["parent_mirror_x"],
"position": track_transform["position"],
"rotation": track_transform["rotation"],
"id": track_id,
"dataIndex": data_index,
**track_shape,
"pos": f"<new Vector3({", ".join([str(x) for x in track_transform["local_position"].tolist()])})>",
"rot": f"<new Quat({", ".join([str(x) for x in track_transform["local_rotation"].as_quat().tolist()])})>",
"connections": connections,
}
@ -170,10 +202,31 @@ def find_tracks(prefab):
switch = try_parse_switch_component(component)
if switch is None:
continue
for track in switch["tracks"]:
tracks.append(parse_track(track))
for data_index, track in enumerate(switch["tracks"]):
tracks.append(parse_track(track, data_index))
add_missing_connections(tracks)
return tracks
def add_missing_connections(tracks):
track_map = {track["id"]: track for track in tracks}
for track in tracks:
for connection in track["connections"]:
if connection["type"] == "<EXTERNAL>":
continue
other_track = track_map.get(connection["internalId"], None)
if other_track is None:
eprint(f"Warning: Track {track['id']} has a connection to an unknown track {connection['internalId']}")
continue
reverse_connections = [conn for conn in other_track["connections"] if conn["type"] == "<INTERNAL>" and conn["internalId"] == track["id"]]
if len(reverse_connections) > 1:
eprint(f"Warning: Track {other_track['id']} has multiple connections to {track["id"]}")
if len(reverse_connections) == 0:
other_track["connections"].append({
"type": "<INTERNAL>",
"end": "<START>" if connection["end"] == "<END>" else "<END>", # Heuristic
"internalId": track["id"]
})
return tracks
def parse_prefab_component(component_yaml, component_type, component_id):
loaded = yaml.load(component_yaml, yaml.SafeLoader)
@ -203,11 +256,15 @@ def resolve_references(value, component_map, visited_ids):
resolve_references(value[key], component_map, visited_ids)
def format_tracks(prefab_name, tracks):
print(f"\"{prefab_name}\": [")
def format_prefab(prefab_name, tracks):
print(f'"{prefab_name}": {{')
print(' "tracks": [')
for track in tracks:
print(f" {json.dumps(track)},")
print("],")
track_object = json.dumps(track)
track_object = re.sub(r'(?<!\\)"<([^<>"]+)>"', r'\1', track_object)
print(f" {track_object},")
print(" ],")
print("},")
def process_file(file_path):
@ -223,7 +280,7 @@ def process_file(file_path):
with open(file_path, encoding="utf8") as infile:
prefab = parse_prefab(infile)
tracks = find_tracks(prefab)
format_tracks(Path(file_path).stem, tracks)
format_prefab(Path(file_path).stem, tracks)
except Exception:
eprint(f"Error processing {file_path}:", file=sys.stderr)
traceback.print_exc(file=sys.stderr)