Преглед изворни кода

Sujith :) ->
1. Added extract TRS functionality

Sujith:) пре 21 часа
родитељ
комит
3235125f3b

+ 9 - 2
Assets/LLM/Editor/DataExtractorMenu.cs

@@ -11,6 +11,7 @@ namespace LLM.Editor
         private const string MidLevelScript = "Assets/LLM/source/orchestrators/extract_mid_level.py";
         private const string LowLevelScript = "Assets/LLM/source/orchestrators/extract_low_level.py";
         private const string CreateConfigScript = "Assets/LLM/source/orchestrators/create_config.py";
+        private const string TRSExtractorScript = "Assets/LLM/source/orchestrators/extract_scene_trs.py";
 
         [MenuItem("Tools/DataExtractor/High Level Export")]
         private static void ExportHighLevel()
@@ -30,6 +31,12 @@ namespace LLM.Editor
             RunExtractor("Low Level", LowLevelScript, true);
         }
         
+        [MenuItem("Tools/DataExtractor/TRS Extractor")]
+        private static void ExportTRS()
+        {
+            RunExtractor("TRS", TRSExtractorScript, true);
+        }
+        
         [MenuItem("Tools/DataExtractor/Create or Update Config File")]
         private static void CreateOrUpdateConfigFile()
         {
@@ -73,7 +80,7 @@ namespace LLM.Editor
 
             var arguments = requireOutputPath ? 
                 $"\"{fullScriptPath}\" --input \"{projectRoot}\" --output \"{outputPath}\"" : 
-                $"{fullScriptPath}";
+                $"\"{fullScriptPath}\"    ";
         
             UnityEngine.Debug.Log($"Running command: \"{pythonExecutable}\" {arguments}");
 
@@ -113,4 +120,4 @@ namespace LLM.Editor
             }
         }
     }
-}
+}

+ 1 - 1
Assets/LLM/source/config/llm_config.json

@@ -1,6 +1,6 @@
 {
     "_comment": "Controls the output of the data extractors. 'shrink_json' (true/false) tokenizes all JSON keys to reduce file size. 'indentation_level' (e.g., 2, 4) controls pretty-printing; use 0 or null for compact, single-line output. `shrink_json` takes precedence on formatting.",
-    "shrink_json": true,
+    "shrink_json": false,
     "indentation_level": 0,
     "ignored_folders": [
         "Assets/LLM",

+ 160 - 0
Assets/LLM/source/orchestrators/extract_scene_trs.py

@@ -0,0 +1,160 @@
+import argparse
+import sys
+import json
+from pathlib import Path
+
+# Add parent directories to the Python path to find utils and parsers
+source_path = Path(__file__).parent.parent
+sys.path.append(str(source_path / 'utils'))
+sys.path.append(str(source_path / 'parsers'))
+
+from file_utils import find_files_by_extension
+from json_utils import write_json
+from config_utils import load_config
+from trs_processor import TRSSceneProcessor
+
+def generate_guid_mappers(input_dir, output_dir, indent=None, shrink=False, ignored_folders=None):
+    """
+    Finds all .meta files and generates JSON files mapping GUIDs to asset paths.
+    This is a necessary setup step to allow the TRS processor to find prefab files.
+    """
+    assets_dir = input_dir / "Assets"
+    if not assets_dir.is_dir():
+        print(f"Error: 'Assets' directory not found in '{input_dir}'", file=sys.stderr)
+        return None
+
+    print("--> Finding all .meta files for GUID mapping...")
+    meta_files = find_files_by_extension(str(assets_dir), '.meta', ignored_folders=ignored_folders, project_root=input_dir)
+    print(f"--> Found {len(meta_files)} .meta files.")
+
+    guid_map = {}
+
+    print("--> Parsing .meta files and mapping GUIDs to asset paths...")
+    for meta_file_path_str in meta_files:
+        meta_file_path = Path(meta_file_path_str)
+        asset_file_path = Path(meta_file_path_str.rsplit('.meta', 1)[0])
+
+        if not asset_file_path.is_file():
+            continue
+
+        guid = None
+        try:
+            with open(meta_file_path, 'r', encoding='utf-8') as f:
+                for line in f:
+                    if line.strip().startswith('guid:'):
+                        guid = line.strip().split(':')[1].strip()
+                        break
+        except Exception as e:
+            print(f"Warning: Could not read or parse guid from {meta_file_path}. {e}", file=sys.stderr)
+            continue
+
+        if guid:
+            # Use the full, resolved path for the guid_map value
+            guid_map[guid] = asset_file_path.resolve().as_posix()
+            
+    print("--> GUID mapping complete.")
+    return guid_map
+
+
+def main():
+    """
+    Main function to run the TRS-only data extraction for scenes.
+    """
+    parser = argparse.ArgumentParser(
+        description="Generates a simplified JSON representation of scene hierarchies with only TRS data."
+    )
+    parser.add_argument(
+        "--input",
+        type=str,
+        required=True,
+        help="The root directory of the target Unity project."
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="The directory where the generated JSON files will be saved."
+    )
+    args = parser.parse_args()
+
+    # --- Load Configuration ---
+    config = load_config()
+    ignored_folders = config.get('ignored_folders', [])
+    shrink_json = config.get('shrink_json', False)
+    indent_level = config.get('indentation_level', 4)
+
+    input_dir = Path(args.input).resolve()
+    output_dir = Path(args.output).resolve()
+
+    if not input_dir.is_dir():
+        print(f"Error: Input path '{input_dir}' is not a valid directory.", file=sys.stderr)
+        sys.exit(1)
+
+    # --- Setup Output Directory ---
+    try:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Output will be saved to: {output_dir}")
+    except OSError as e:
+        print(f"Error: Could not create output directory '{output_dir}'. {e}", file=sys.stderr)
+        sys.exit(1)
+
+    assets_dir = input_dir / "Assets"
+    if not assets_dir.is_dir():
+        print(f"Warning: 'Assets' directory not found in '{input_dir}'. Skipping all processing.", file=sys.stderr)
+        return
+
+    print("\n--- Running Scene TRS Extraction ---")
+
+    # --- Task 1: Generate GUID Map (for prefab lookups) ---
+    print("\n[1/2] Generating GUID Mappers...")
+    guid_map = generate_guid_mappers(
+        input_dir, 
+        output_dir, # Not saving mappers, just using the map in memory
+        ignored_folders=ignored_folders
+    )
+    if guid_map is None:
+        print("Error: Failed to generate GUID map. Aborting.", file=sys.stderr)
+        sys.exit(1)
+    print("--> GUID Mapper generation complete.")
+
+    # --- Task 2: Find and Process Scene Files ---
+    print("\n[2/2] Parsing Scene Hierarchies for TRS data...")
+    
+    scene_files = find_files_by_extension(str(assets_dir), '.unity', ignored_folders=ignored_folders, project_root=input_dir)
+    
+    if not scene_files:
+        print("--> No scene files found to process.")
+        return
+        
+    print(f"--> Found {len(scene_files)} scene file(s) to process.")
+
+    total_files = len(scene_files)
+    for i, file_path_str in enumerate(scene_files):
+        file_path = Path(file_path_str)
+        
+        output_json_path = output_dir / f"{file_path.stem}.json"
+        
+        try:
+            print(f"\n--- Processing {file_path.name} ({i+1}/{total_files}) ---")
+            
+            processor = TRSSceneProcessor(guid_map)
+            
+            print(f"    -> Loading and parsing file...")
+            if not processor.load_documents(file_path):
+                print(f"Warning: Could not load or parse {file_path.name}. Skipping.", file=sys.stderr)
+                continue
+
+            print(f"    -> Building hierarchy and extracting TRS data...")
+            hierarchy = processor.process()
+            
+            # The hierarchy is the list of root objects, which will be the top-level JSON array.
+            write_json(hierarchy, output_json_path, indent=indent_level, shrink=shrink_json)
+            print(f"--> Successfully processed TRS hierarchy for {file_path.name} -> {output_json_path}")
+
+        except Exception as e:
+            print(f"Error processing TRS hierarchy for {file_path.name}: {e}", file=sys.stderr)
+
+    print("\nScene TRS extraction complete.")
+
+if __name__ == "__main__":
+    main()

+ 7 - 0
Assets/LLM/source/orchestrators/extract_scene_trs.py.meta

@@ -0,0 +1,7 @@
+fileFormatVersion: 2
+guid: 47ba9a6823c44481ca23ed7363ca6d70
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

BIN
Assets/LLM/source/parsers/__pycache__/trs_processor.cpython-313.pyc


+ 7 - 0
Assets/LLM/source/parsers/__pycache__/trs_processor.cpython-313.pyc.meta

@@ -0,0 +1,7 @@
+fileFormatVersion: 2
+guid: 22c27d35daf4744faa0571fc5359e337
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 

+ 260 - 0
Assets/LLM/source/parsers/trs_processor.py

@@ -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

+ 7 - 0
Assets/LLM/source/parsers/trs_processor.py.meta

@@ -0,0 +1,7 @@
+fileFormatVersion: 2
+guid: 703dfbb0286a04baf9fec825d7f81840
+DefaultImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: