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