import argparse import sys import json from pathlib import Path # Add the utils directory to the Python path utils_path = Path(__file__).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 UnitySceneProcessor: def __init__(self, guid_map): self.guid_map = guid_map self.object_map = {} self.nodes = {} self.prefab_nodes = {} self.transform_to_gameobject = {} self.gameobject_to_transform = {} self.transform_children = {} self.stripped_gameobjects = {} self.processed_relationships = set() def load_prefab_data(self, prefab_guid): """Load and parse prefab data from GUID""" if not prefab_guid or prefab_guid not in self.guid_map: return {} prefab_path = self.guid_map[prefab_guid] try: documents = load_unity_yaml(prefab_path) if not documents: return {} raw_object_map = {int(doc.anchor.value): doc for doc in documents if hasattr(doc, 'anchor') and doc.anchor is not None} return {file_id: convert_to_plain_python_types(obj) for file_id, obj in raw_object_map.items()} except Exception: return {} def calculate_deep_sibling_index(self, scene_transform_id, prefab_guid): """Calculate deep sibling index for scene objects that are children of prefab objects""" scene_transform = self.object_map.get(scene_transform_id, {}).get('Transform', {}) scene_root_order = scene_transform.get('m_RootOrder', 0) # Get parent transform (should be stripped) parent_transform_id = scene_transform.get('m_Father', {}).get('fileID') if not parent_transform_id: return str(scene_root_order) parent_transform = self.object_map.get(parent_transform_id, {}).get('Transform', {}) corresponding_source = parent_transform.get('m_CorrespondingSourceObject', {}) prefab_transform_id = corresponding_source.get('fileID') if not prefab_transform_id: return str(scene_root_order) # Load prefab data and traverse hierarchy prefab_data = self.load_prefab_data(prefab_guid) if not prefab_data: return str(scene_root_order) # Build sibling index by traversing up the prefab hierarchy sibling_indices = [] current_transform_id = prefab_transform_id while current_transform_id and current_transform_id in prefab_data: transform_data = prefab_data.get(current_transform_id, {}).get('Transform', {}) if not transform_data: break root_order = transform_data.get('m_RootOrder', 0) sibling_indices.insert(0, str(root_order)) # Move to parent parent_id = transform_data.get('m_Father', {}).get('fileID') if not parent_id or parent_id == 0: break current_transform_id = parent_id # Add the scene object's own root order at the end sibling_indices.append(str(scene_root_order)) return '-'.join(sibling_indices) def process_first_pass(self): """First pass: Build relationship maps and create basic nodes""" stripped_transforms = {} stripped_gameobjects = {} prefab_instances = {} for file_id, obj_data in self.object_map.items(): if 'GameObject' in obj_data: go_info = obj_data['GameObject'] # Always create a node for a GameObject. # If it's part of a prefab, we'll link it later. self.nodes[file_id] = { 'fileID': str(file_id), 'm_Name': go_info.get('m_Name', 'Unknown'), 'm_IsActive': go_info.get('m_IsActive', 1), 'm_TagString': go_info.get('m_TagString', 'Untagged'), 'm_Layer': go_info.get('m_Layer', 0), 'components': [], 'children': [] } # If it's a stripped GameObject, track it for component linking is_stripped = any('stripped' in str(key) for key in obj_data.keys() if hasattr(key, '__str__')) if is_stripped: prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID') if prefab_instance_id: stripped_gameobjects[file_id] = { 'prefab_instance_id': prefab_instance_id, 'm_CorrespondingSourceObject': go_info.get('m_CorrespondingSourceObject', {}) } # If it's the root GameObject of a prefab instance, add a link for merging prefab_instance_id = go_info.get('m_PrefabInstance', {}).get('fileID') if prefab_instance_id and not is_stripped: self.nodes[file_id]['prefab_instance_id'] = prefab_instance_id elif 'PrefabInstance' in obj_data: prefab_info = obj_data['PrefabInstance'] source_prefab = prefab_info.get('m_SourcePrefab', {}) modifications = prefab_info.get('m_Modification', {}) # Extract m_IsActive, m_Name, and m_RootOrder from modifications m_is_active = 1 # default m_name = None m_root_order = 999999 # default for mod in modifications.get('m_Modifications', []): if mod.get('propertyPath') == 'm_IsActive': m_is_active = mod.get('value', 1) elif mod.get('propertyPath') == 'm_Name': m_name = mod.get('value') elif mod.get('propertyPath') == 'm_RootOrder': m_root_order = mod.get('value', 999999) prefab_instances[file_id] = { 'source_prefab_guid': source_prefab.get('guid'), 'm_TransformParent': modifications.get('m_TransformParent', {}).get('fileID'), 'm_IsActive': m_is_active, 'm_Name': m_name, 'm_RootOrder': m_root_order } elif 'Transform' in obj_data: transform_info = obj_data['Transform'] # Check if this is a stripped transform (has m_PrefabInstance) prefab_instance_id = transform_info.get('m_PrefabInstance', {}).get('fileID') if prefab_instance_id: stripped_transforms[file_id] = { 'prefab_instance_id': prefab_instance_id, 'm_CorrespondingSourceObject': transform_info.get('m_CorrespondingSourceObject', {}) } # Build transform mappings 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 # Build parent-child relationships 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) # Create prefab nodes for prefab_id, prefab_info in prefab_instances.items(): self.prefab_nodes[prefab_id] = { 'fileID': str(prefab_id), 'source_prefab_guid': prefab_info['source_prefab_guid'], 'm_IsActive': prefab_info['m_IsActive'], 'm_RootOrder': prefab_info.get('m_RootOrder', 999999), 'components': [], # Will be populated with override components 'children': [] } if prefab_info.get('m_Name'): self.prefab_nodes[prefab_id]['m_Name'] = prefab_info['m_Name'] # Store for second pass self.stripped_transforms = stripped_transforms self.stripped_gameobjects = stripped_gameobjects self.prefab_instances = prefab_instances def process_second_pass(self): """Second pass: Build hierarchy relationships""" # 1. Handle standard GameObject parent-child relationships for go_id, node in self.nodes.items(): transform_id = self.gameobject_to_transform.get(go_id) if not transform_id: continue # Get child transforms child_transform_ids = self.transform_children.get(transform_id, []) for child_transform_id in child_transform_ids: # Check if child transform is stripped (part of prefab) if child_transform_id in self.stripped_transforms: stripped_info = self.stripped_transforms[child_transform_id] prefab_instance_id = stripped_info['prefab_instance_id'] if prefab_instance_id in self.prefab_nodes: relationship_key = f"{go_id}->{prefab_instance_id}" if relationship_key not in self.processed_relationships: node['children'].append(self.prefab_nodes[prefab_instance_id]) self.processed_relationships.add(relationship_key) else: # Regular GameObject child child_go_id = self.transform_to_gameobject.get(child_transform_id) if child_go_id and child_go_id in self.nodes: relationship_key = f"{go_id}->{child_go_id}" if relationship_key not in self.processed_relationships: node['children'].append(self.nodes[child_go_id]) self.processed_relationships.add(relationship_key) # 2. Handle prefab-to-parent relationships for prefab_id, prefab_info in self.prefab_instances.items(): parent_transform_id = prefab_info['m_TransformParent'] if not parent_transform_id or parent_transform_id == 0: continue # Root prefab, will be handled in get_root_objects # Find parent GameObject or parent prefab parent_go_id = self.transform_to_gameobject.get(parent_transform_id) if parent_go_id and parent_go_id in self.nodes: # Prefab child of GameObject relationship_key = f"{parent_go_id}->{prefab_id}" if relationship_key not in self.processed_relationships: self.nodes[parent_go_id]['children'].append(self.prefab_nodes[prefab_id]) self.processed_relationships.add(relationship_key) elif parent_transform_id in self.stripped_transforms: # Prefab child of another prefab stripped_info = self.stripped_transforms[parent_transform_id] parent_prefab_id = stripped_info['prefab_instance_id'] if parent_prefab_id in self.prefab_nodes: relationship_key = f"{parent_prefab_id}->{prefab_id}" if relationship_key not in self.processed_relationships: self.prefab_nodes[parent_prefab_id]['children'].append(self.prefab_nodes[prefab_id]) self.processed_relationships.add(relationship_key) # 3. Handle scene objects as children of prefabs (with deep sibling index) for go_id, node in self.nodes.items(): transform_id = self.gameobject_to_transform.get(go_id) if not transform_id: continue # Find parent transform parent_transform_id = None for parent_id, children in self.transform_children.items(): if transform_id in children: parent_transform_id = parent_id break if parent_transform_id and parent_transform_id in self.stripped_transforms: stripped_info = self.stripped_transforms[parent_transform_id] prefab_instance_id = stripped_info['prefab_instance_id'] if prefab_instance_id in self.prefab_nodes: # Calculate deep sibling index prefab_guid = self.prefab_instances[prefab_instance_id]['source_prefab_guid'] deep_sibling_index = self.calculate_deep_sibling_index(transform_id, prefab_guid) if deep_sibling_index: node['deep_sibling_index'] = deep_sibling_index # Add as child of prefab relationship_key = f"{prefab_instance_id}->{go_id}" if relationship_key not in self.processed_relationships: self.prefab_nodes[prefab_instance_id]['children'].append(node) self.processed_relationships.add(relationship_key) def process_third_pass(self): """Third pass: Extract components for both GameObjects and PrefabInstances""" for file_id, obj_data in self.object_map.items(): # Skip GameObjects, Transforms, and PrefabInstances if any(key in obj_data for key in ['GameObject', 'Transform', 'PrefabInstance']): continue # Find the GameObject this component belongs to gameobject_id = None for key, data in obj_data.items(): if isinstance(data, dict) and 'm_GameObject' in data: gameobject_id = data['m_GameObject'].get('fileID') break if not gameobject_id: continue component_name = next((key for key in obj_data if key not in ['m_ObjectHideFlags']), 'Unknown') # Handle MonoBehaviour with script GUID if component_name == 'MonoBehaviour': script_info = obj_data['MonoBehaviour'].get('m_Script', {}) script_guid = script_info.get('guid') if script_guid: component_entry = script_guid else: component_entry = "MonoBehaviour" else: component_entry = component_name # Check if this component belongs to a regular GameObject if gameobject_id in self.nodes: if component_entry not in self.nodes[gameobject_id]['components']: self.nodes[gameobject_id]['components'].append(component_entry) # Check if this component belongs to a stripped GameObject (prefab override) elif gameobject_id in self.stripped_gameobjects: stripped_info = self.stripped_gameobjects[gameobject_id] prefab_instance_id = stripped_info['prefab_instance_id'] if prefab_instance_id in self.prefab_nodes: if component_entry not in self.prefab_nodes[prefab_instance_id]['components']: self.prefab_nodes[prefab_instance_id]['components'].append(component_entry) # If neither, this component belongs to an unknown GameObject (potential orphan) else: print(f"Warning: Component {component_name} (ID: {file_id}) references unknown GameObject {gameobject_id}") def verification_pass(self): """ Verifies and fixes parent-child relationships that might have been missed by using a direct child->parent lookup. This acts as a patch for the less reliable reverse-lookup used in process_second_pass. """ for go_id, node in list(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 not parent_transform_id or parent_transform_id == 0: continue # This is a root object, no parent to verify. # Case 1: Parent is a PrefabInstance (via a stripped transform) if parent_transform_id in self.stripped_transforms: stripped_info = self.stripped_transforms[parent_transform_id] parent_prefab_id = stripped_info['prefab_instance_id'] if parent_prefab_id in self.prefab_nodes: relationship_key = f"{parent_prefab_id}->{go_id}" if relationship_key not in self.processed_relationships: # Add deep sibling index if it's a child of a prefab prefab_guid = self.prefab_instances[parent_prefab_id]['source_prefab_guid'] deep_sibling_index = self.calculate_deep_sibling_index(transform_id, prefab_guid) if deep_sibling_index: node['deep_sibling_index'] = deep_sibling_index self.prefab_nodes[parent_prefab_id]['children'].append(node) self.processed_relationships.add(relationship_key) # Case 2: Parent is a regular GameObject else: parent_go_id = self.transform_to_gameobject.get(parent_transform_id) if parent_go_id and parent_go_id in self.nodes: relationship_key = f"{parent_go_id}->{go_id}" if relationship_key not in self.processed_relationships: self.nodes[parent_go_id]['children'].append(node) self.processed_relationships.add(relationship_key) def merge_prefab_data_pass(self): """Fourth pass: Merge GameObject data into their corresponding PrefabInstance nodes""" nodes_to_delete = [] for go_id, go_node in self.nodes.items(): prefab_instance_id = go_node.get('prefab_instance_id') if not prefab_instance_id: continue if prefab_instance_id in self.prefab_nodes: # Merge the GameObject data into the PrefabInstance node # The prefab node's existing data takes precedence prefab_node = self.prefab_nodes[prefab_instance_id] # Create a copy of the gameobject node to avoid modifying during iteration go_node_copy = go_node.copy() # Remove keys that should not be merged or are already handled del go_node_copy['fileID'] del go_node_copy['prefab_instance_id'] # Update prefab_node with go_node_copy data # This will overwrite keys in prefab_node with values from go_node_copy # if they exist in both. We want the opposite. # Let's do it manually to ensure prefab data is kept for key, value in go_node_copy.items(): if key not in prefab_node: prefab_node[key] = value # Specifically handle children and components to merge them prefab_node['children'].extend(go_node_copy.get('children', [])) # Merge components, avoiding duplicates existing_components = set(str(c) for c in prefab_node.get('components', [])) for comp in go_node_copy.get('components', []): if str(comp) not in existing_components: prefab_node['components'].append(comp) # Rename 'components' to 'addedComponents' for clarity if 'components' in prefab_node: prefab_node['addedComponents'] = prefab_node.pop('components') # Mark the original GameObject node for deletion nodes_to_delete.append(go_id) # Remove the now-redundant GameObject nodes for go_id in nodes_to_delete: del self.nodes[go_id] def cleanup_pass(self, nodes): """ Recursively cleans up temporary or internal properties from the final node structure. """ # List of properties to remove from the final output cleanup_keys = ['m_RootOrder'] for node in nodes: for key in cleanup_keys: if key in node: del node[key] if 'children' in node and node['children']: self.cleanup_pass(node['children']) def process_file(self, file_path): """Main processing method""" # Load and parse the file documents = load_unity_yaml(file_path) if not documents: return [] # Build object map 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()} # Process in passes self.process_first_pass() self.process_second_pass() self.verification_pass() self.process_third_pass() self.merge_prefab_data_pass() # Use the centralized parser to get the final, sorted root objects parser = HierarchyParser(self.object_map) root_object_ids = parser.get_root_object_ids() # Build the final list of nodes from the identified roots root_nodes = [] all_nodes = {**self.nodes, **self.prefab_nodes} for file_id, _ in root_object_ids: if file_id in all_nodes: root_nodes.append(all_nodes[file_id]) # Run the final cleanup pass self.cleanup_pass(root_nodes) return root_nodes if __name__ == "__main__": # This script is intended to be used as a module. # For direct execution, see extract_mid_level.py. pass