|
@@ -0,0 +1,260 @@
|
|
|
+import sys
|
|
|
+from pathlib import Path
|
|
|
+
|
|
|
+# Add the parent directory to the Python path to find utils
|
|
|
+utils_path = Path(__file__).parent.parent / 'utils'
|
|
|
+sys.path.append(str(utils_path))
|
|
|
+
|
|
|
+from yaml_utils import load_unity_yaml, convert_to_plain_python_types
|
|
|
+from hierarchy_utils import HierarchyParser
|
|
|
+
|
|
|
+class TRSSceneProcessor:
|
|
|
+ """
|
|
|
+ A specialized processor that uses the robust multi-pass system from the mid-level
|
|
|
+ extractor to build a correct scene hierarchy, then extracts only TRS data.
|
|
|
+ """
|
|
|
+ def __init__(self, guid_map):
|
|
|
+ self.guid_map = guid_map
|
|
|
+ self.object_map = {}
|
|
|
+ self.prefab_cache = {}
|
|
|
+ self.nodes = {}
|
|
|
+ self.prefab_nodes = {}
|
|
|
+ self.transform_to_gameobject = {}
|
|
|
+ self.gameobject_to_transform = {}
|
|
|
+ self.transform_children = {}
|
|
|
+ self.stripped_gameobjects = {}
|
|
|
+ self.stripped_transforms = {}
|
|
|
+ self.prefab_instances = {}
|
|
|
+ self.processed_relationships = set()
|
|
|
+
|
|
|
+ def load_documents(self, file_path):
|
|
|
+ """Loads and parses the yaml documents from a scene/prefab file."""
|
|
|
+ documents = load_unity_yaml(file_path)
|
|
|
+ if not documents:
|
|
|
+ self.object_map = {}
|
|
|
+ return False
|
|
|
+
|
|
|
+ raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
|
|
|
+ self.object_map = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}
|
|
|
+ return True
|
|
|
+
|
|
|
+ def _load_prefab_data(self, prefab_guid):
|
|
|
+ """Load and cache prefab data from its GUID."""
|
|
|
+ if prefab_guid in self.prefab_cache:
|
|
|
+ return self.prefab_cache[prefab_guid]
|
|
|
+
|
|
|
+ if not prefab_guid or prefab_guid not in self.guid_map:
|
|
|
+ self.prefab_cache[prefab_guid] = None
|
|
|
+ return None
|
|
|
+
|
|
|
+ prefab_path = self.guid_map[prefab_guid]
|
|
|
+ try:
|
|
|
+ documents = load_unity_yaml(prefab_path)
|
|
|
+ if not documents:
|
|
|
+ self.prefab_cache[prefab_guid] = None
|
|
|
+ return None
|
|
|
+
|
|
|
+ raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None}
|
|
|
+ prefab_data = {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()}
|
|
|
+ self.prefab_cache[prefab_guid] = prefab_data
|
|
|
+ return prefab_data
|
|
|
+ except Exception:
|
|
|
+ self.prefab_cache[prefab_guid] = None
|
|
|
+ return None
|
|
|
+
|
|
|
+ def build_relationship_maps(self):
|
|
|
+ """Pass 1: Build all necessary maps for hierarchy construction."""
|
|
|
+ for file_id, obj_data in self.object_map.items():
|
|
|
+ if 'GameObject' in obj_data:
|
|
|
+ go_info = obj_data['GameObject']
|
|
|
+ self.nodes[file_id] = {'fileID': str(file_id), 'm_Name': go_info.get('m_Name', 'Unknown'), 'children': []}
|
|
|
+ if any('stripped' in str(k) for k in obj_data.keys()):
|
|
|
+ self.stripped_gameobjects[file_id] = go_info
|
|
|
+ prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID')
|
|
|
+ if prefab_instance_id and not file_id in self.stripped_gameobjects:
|
|
|
+ self.nodes[file_id]['prefab_instance_id'] = prefab_instance_id
|
|
|
+
|
|
|
+ elif 'PrefabInstance' in obj_data:
|
|
|
+ prefab_info = obj_data['PrefabInstance']
|
|
|
+ self.prefab_instances[file_id] = prefab_info
|
|
|
+
|
|
|
+ # Get name from modifications or fallback to prefab asset filename
|
|
|
+ mod_name = next((mod['value'] for mod in prefab_info.get('m_Modification', {}).get('m_Modifications', []) if mod.get('propertyPath') == 'm_Name'), None)
|
|
|
+
|
|
|
+ if not mod_name:
|
|
|
+ guid = prefab_info.get('m_SourcePrefab', {}).get('guid')
|
|
|
+ if guid and guid in self.guid_map:
|
|
|
+ prefab_path = Path(self.guid_map[guid])
|
|
|
+ mod_name = prefab_path.stem
|
|
|
+ else:
|
|
|
+ mod_name = 'Unknown'
|
|
|
+
|
|
|
+ self.prefab_nodes[file_id] = {'fileID': str(file_id), 'm_Name': mod_name, 'children': []}
|
|
|
+
|
|
|
+ elif 'Transform' in obj_data:
|
|
|
+ transform_info = obj_data['Transform']
|
|
|
+ gameobject_id = transform_info.get('m_GameObject', {}).get('fileID')
|
|
|
+ if gameobject_id:
|
|
|
+ self.transform_to_gameobject[file_id] = gameobject_id
|
|
|
+ self.gameobject_to_transform[gameobject_id] = file_id
|
|
|
+
|
|
|
+ parent_id = transform_info.get('m_Father', {}).get('fileID')
|
|
|
+ if parent_id and parent_id != 0:
|
|
|
+ if parent_id not in self.transform_children:
|
|
|
+ self.transform_children[parent_id] = []
|
|
|
+ self.transform_children[parent_id].append(file_id)
|
|
|
+
|
|
|
+ if transform_info.get('m_PrefabInstance', {}).get('fileID'):
|
|
|
+ self.stripped_transforms[file_id] = transform_info
|
|
|
+
|
|
|
+ def build_hierarchy(self):
|
|
|
+ """Pass 2: Connect nodes into a hierarchy."""
|
|
|
+ # Link standard GameObjects
|
|
|
+ for go_id, node in self.nodes.items():
|
|
|
+ transform_id = self.gameobject_to_transform.get(go_id)
|
|
|
+ if not transform_id: continue
|
|
|
+
|
|
|
+ child_transform_ids = self.transform_children.get(transform_id, [])
|
|
|
+ for child_transform_id in child_transform_ids:
|
|
|
+ child_go_id = self.transform_to_gameobject.get(child_transform_id)
|
|
|
+ if child_go_id and child_go_id in self.nodes:
|
|
|
+ node['children'].append(self.nodes[child_go_id])
|
|
|
+
|
|
|
+ # Link prefab instances to their parents
|
|
|
+ for prefab_id, prefab_info in self.prefab_instances.items():
|
|
|
+ parent_transform_id = prefab_info.get('m_Modification', {}).get('m_TransformParent', {}).get('fileID')
|
|
|
+ if not parent_transform_id: continue
|
|
|
+
|
|
|
+ parent_go_id = self.transform_to_gameobject.get(parent_transform_id)
|
|
|
+ if parent_go_id and parent_go_id in self.nodes:
|
|
|
+ self.nodes[parent_go_id]['children'].append(self.prefab_nodes[prefab_id])
|
|
|
+ elif parent_transform_id in self.stripped_transforms:
|
|
|
+ parent_prefab_instance_id = self.stripped_transforms[parent_transform_id].get('m_PrefabInstance', {}).get('fileID')
|
|
|
+ if parent_prefab_instance_id in self.prefab_nodes:
|
|
|
+ self.prefab_nodes[parent_prefab_instance_id]['children'].append(self.prefab_nodes[prefab_id])
|
|
|
+
|
|
|
+ # Link scene objects that are children of prefabs
|
|
|
+ for go_id, node in self.nodes.items():
|
|
|
+ transform_id = self.gameobject_to_transform.get(go_id)
|
|
|
+ if not transform_id: continue
|
|
|
+ transform_info = self.object_map.get(transform_id, {}).get('Transform', {})
|
|
|
+ parent_transform_id = transform_info.get('m_Father', {}).get('fileID')
|
|
|
+ if parent_transform_id in self.stripped_transforms:
|
|
|
+ parent_prefab_instance_id = self.stripped_transforms[parent_transform_id].get('m_PrefabInstance', {}).get('fileID')
|
|
|
+ if parent_prefab_instance_id in self.prefab_nodes:
|
|
|
+ self.prefab_nodes[parent_prefab_instance_id]['children'].append(node)
|
|
|
+
|
|
|
+ def merge_prefab_info(self):
|
|
|
+ """Pass 3: Merge root GameObject data into PrefabInstance nodes."""
|
|
|
+ nodes_to_delete = []
|
|
|
+ for go_id, go_node in self.nodes.items():
|
|
|
+ prefab_instance_id = go_node.get('prefab_instance_id')
|
|
|
+ if prefab_instance_id and prefab_instance_id in self.prefab_nodes:
|
|
|
+ prefab_node = self.prefab_nodes[prefab_instance_id]
|
|
|
+ prefab_info = self.prefab_instances[prefab_instance_id]
|
|
|
+
|
|
|
+ # Name is now set in build_relationship_maps. We just merge the rest.
|
|
|
+ prefab_node['source_prefab_guid'] = prefab_info.get('m_SourcePrefab', {}).get('guid')
|
|
|
+ prefab_node['children'].extend(go_node.get('children', []))
|
|
|
+ nodes_to_delete.append(go_id)
|
|
|
+
|
|
|
+ for go_id in nodes_to_delete:
|
|
|
+ del self.nodes[go_id]
|
|
|
+
|
|
|
+ def extract_trs_data(self, nodes):
|
|
|
+ """Pass 4: Recursively find and attach TRS data to each node."""
|
|
|
+ for node in nodes:
|
|
|
+ file_id = int(node['fileID'])
|
|
|
+
|
|
|
+ # Default TRS values
|
|
|
+ trs = {"position": {"x": 0, "y": 0, "z": 0}, "rotation": {"x": 0, "y": 0, "z": 0, "w": 1}, "scale": {"x": 1, "y": 1, "z": 1}}
|
|
|
+
|
|
|
+ if 'source_prefab_guid' in node: # It's a prefab instance
|
|
|
+ prefab_instance_info = self.prefab_instances.get(file_id, {})
|
|
|
+ modifications = prefab_instance_info.get('m_Modification', {}).get('m_Modifications', [])
|
|
|
+
|
|
|
+ # Find the root transform of the prefab to get base TRS
|
|
|
+ prefab_data = self._load_prefab_data(node['source_prefab_guid'])
|
|
|
+ if prefab_data:
|
|
|
+ for obj in prefab_data.values():
|
|
|
+ if 'Transform' in obj and obj['Transform'].get('m_Father', {}).get('fileID') == 0:
|
|
|
+ base_transform = obj['Transform']
|
|
|
+ trs['position'] = base_transform.get('m_LocalPosition', trs['position'])
|
|
|
+ trs['rotation'] = base_transform.get('m_LocalRotation', trs['rotation'])
|
|
|
+ trs['scale'] = base_transform.get('m_LocalScale', trs['scale'])
|
|
|
+ break
|
|
|
+
|
|
|
+ # Apply modifications
|
|
|
+ for mod in modifications:
|
|
|
+ prop = mod.get('propertyPath', '')
|
|
|
+
|
|
|
+ # Skip any property that is not part of a Transform
|
|
|
+ if not ('m_LocalPosition' in prop or 'm_LocalRotation' in prop or 'm_LocalScale' in prop):
|
|
|
+ continue
|
|
|
+
|
|
|
+ try:
|
|
|
+ val = float(mod.get('value', 0))
|
|
|
+ except (ValueError, TypeError):
|
|
|
+ # If the value is not a valid float, skip this modification
|
|
|
+ print(f"Warning: Could not convert value '{mod.get('value')}' for property '{prop}' to float. Skipping.", file=sys.stderr)
|
|
|
+ continue
|
|
|
+
|
|
|
+ if prop == 'm_LocalPosition.x': trs['position']['x'] = val
|
|
|
+ elif prop == 'm_LocalPosition.y': trs['position']['y'] = val
|
|
|
+ elif prop == 'm_LocalPosition.z': trs['position']['z'] = val
|
|
|
+ elif prop == 'm_LocalRotation.x': trs['rotation']['x'] = val
|
|
|
+ elif prop == 'm_LocalRotation.y': trs['rotation']['y'] = val
|
|
|
+ elif prop == 'm_LocalRotation.z': trs['rotation']['z'] = val
|
|
|
+ elif prop == 'm_LocalRotation.w': trs['rotation']['w'] = val
|
|
|
+ elif prop == 'm_LocalScale.x': trs['scale']['x'] = val
|
|
|
+ elif prop == 'm_LocalScale.y': trs['scale']['y'] = val
|
|
|
+ elif prop == 'm_LocalScale.z': trs['scale']['z'] = val
|
|
|
+
|
|
|
+ else: # It's a regular GameObject
|
|
|
+ transform_id = self.gameobject_to_transform.get(file_id)
|
|
|
+ if transform_id and transform_id in self.object_map:
|
|
|
+ transform_info = self.object_map[transform_id].get('Transform', {})
|
|
|
+ trs['position'] = transform_info.get('m_LocalPosition', trs['position'])
|
|
|
+ trs['rotation'] = transform_info.get('m_LocalRotation', trs['rotation'])
|
|
|
+ trs['scale'] = transform_info.get('m_LocalScale', trs['scale'])
|
|
|
+
|
|
|
+ node['trs'] = trs
|
|
|
+
|
|
|
+ if 'children' in node and node['children']:
|
|
|
+ self.extract_trs_data(node['children'])
|
|
|
+
|
|
|
+ def cleanup_hierarchy(self, nodes):
|
|
|
+ """Pass 5: Remove intermediate data, leaving only the desired fields."""
|
|
|
+ final_nodes = []
|
|
|
+ for node in nodes:
|
|
|
+ clean_node = {
|
|
|
+ 'name': node.get('m_Name', 'Unknown'),
|
|
|
+ 'trs': node.get('trs', {}),
|
|
|
+ 'children': self.cleanup_hierarchy(node.get('children', []))
|
|
|
+ }
|
|
|
+ if 'source_prefab_guid' in node:
|
|
|
+ clean_node['source_prefab_guid'] = node['source_prefab_guid']
|
|
|
+ final_nodes.append(clean_node)
|
|
|
+ return final_nodes
|
|
|
+
|
|
|
+ def process(self):
|
|
|
+ """Main processing pipeline."""
|
|
|
+ if not self.object_map:
|
|
|
+ return []
|
|
|
+
|
|
|
+ # --- Run all passes ---
|
|
|
+ self.build_relationship_maps()
|
|
|
+ self.build_hierarchy()
|
|
|
+ self.merge_prefab_info()
|
|
|
+
|
|
|
+ # Get root nodes for the final passes
|
|
|
+ parser = HierarchyParser(self.object_map)
|
|
|
+ root_object_ids = parser.get_root_object_ids()
|
|
|
+
|
|
|
+ all_nodes = {**self.nodes, **self.prefab_nodes}
|
|
|
+ root_nodes = [all_nodes[file_id] for file_id, _ in root_object_ids if file_id in all_nodes]
|
|
|
+
|
|
|
+ self.extract_trs_data(root_nodes)
|
|
|
+ final_hierarchy = self.cleanup_hierarchy(root_nodes)
|
|
|
+
|
|
|
+ return final_hierarchy
|