From 0771cf47b500bb47034033aea682016c5fe1a951 Mon Sep 17 00:00:00 2001 From: Jan Vopicka Date: Tue, 27 Aug 2024 14:22:52 +0200 Subject: [PATCH] feat(loader): Option to preserve scene hierarchy when loading blend --- blenderproc/python/loader/BlendLoader.py | 81 ++++++++++++++++++++++-- 1 file changed, 76 insertions(+), 5 deletions(-) diff --git a/blenderproc/python/loader/BlendLoader.py b/blenderproc/python/loader/BlendLoader.py index b7cd371a0..13ff75ced 100644 --- a/blenderproc/python/loader/BlendLoader.py +++ b/blenderproc/python/loader/BlendLoader.py @@ -1,17 +1,79 @@ """Loading the content of .blend files""" import re -from typing import List, Union, Optional +import os +import collections +from typing import List, Dict, Tuple, Union, Optional import bpy from blenderproc.python.utility.BlenderUtility import collect_all_orphan_data_blocks from blenderproc.python.types.EntityUtility import Entity, convert_to_entity_subclass -from blenderproc.python.utility.Utility import resolve_path +from blenderproc.python.utility.Utility import Utility, resolve_path +def record_hierarchy(path: str) -> Tuple[Dict[str, str], Dict[str, str]]: + """ + Backup already set up scene, open blend file, store it's hierarchy and restore original scene. + + Hierarchy cannot be read by just loading the library, the scene needs to be open to access it. + This approach is easier than resetting up already set up scene. + :param path: Path to a .blend file. + :return: Tuple of two dictionaries, one with collection and list of its children and one with + object and list of collections it belongs to. + """ + temp_filepath = os.path.join(Utility.get_temporary_directory(), "temp.blend") + bpy.ops.wm.save_as_mainfile(filepath=temp_filepath) + bpy.ops.wm.open_mainfile(filepath=path) + + collection_hierarchy = {} + object_collections_hierarchy = collections.defaultdict(list) + + for coll in bpy.data.collections: + collection_hierarchy[coll.name] = [c.name for c in coll.children] + for obj in coll.objects: + object_collections_hierarchy[obj.name].append(coll.name) + + bpy.ops.wm.open_mainfile(filepath=temp_filepath) + os.remove(temp_filepath) + + return collection_hierarchy, object_collections_hierarchy + +def build_collection_hierarchy(collection_hierarchy: Dict[str, List[str]]): + """ + Recreate collection hierarchy from source blend + + :param collection_hierarchy: Dictionary of collection names and list oif their children names + """ + + root = bpy.context.scene.collection + for coll in collection_hierarchy: + new_coll = bpy.data.collections.new(coll) + root.children.link(new_coll) + + for coll, children in collection_hierarchy.items(): + for child in children: + child_coll = bpy.data.collections.get(child) + bpy.data.collections.get(coll).children.link(child_coll) + root.children.unlink(child_coll) + +def add_obj_to_hierarchy(obj: bpy.types.Object, object_collections_hierarchy: Dict[str, List[str]]): + """ + Add object to collection of same name as the one in source blend + + :param collection_hierarchy: Dictionary of collection names and list oif their children names + """ + collection_parents = object_collections_hierarchy.get(obj.name) + if collection_parents is None: + bpy.context.collection.objects.link(obj) + else: + for parent in collection_parents: + coll = bpy.data.collections.get(parent) + assert coll is not None + coll.objects.link(obj) + def load_blend(path: str, obj_types: Optional[Union[List[str], str]] = None, name_regrex: Optional[str] = None, - data_blocks: Union[List[str], str] = "objects", link: bool = False) -> List[Entity]: + data_blocks: Union[List[str], str] = "objects", link: bool = False, store_hierarchy: bool = False) -> List[Entity]: """ Loads entities (everything that can be stored in a .blend file's folders, see Blender's documentation for bpy.types.ID for more info) that match a name pattern from a specified .blend file's section/data_block. @@ -27,6 +89,7 @@ def load_blend(path: str, obj_types: Optional[Union[List[str], str]] = None, nam Available options are: ['armatures', 'cameras', 'curves', 'hairs', 'hair_curves', 'images', 'lights', 'materials', 'meshes', 'objects', 'textures'] :param link: whether to link instead of append data blocks from .blend file. Linked objects can not be modified. + :param store_hierarchy: Maintain top-level hierarchy from scene provided by path param. :return: The list of loaded mesh objects. """ if obj_types is None: @@ -41,6 +104,9 @@ def load_blend(path: str, obj_types: Optional[Union[List[str], str]] = None, nam # Remember which orphans existed beforehand orphans_before = collect_all_orphan_data_blocks() + if store_hierarchy: + collection_hierarchy, object_collections_hierarchy = record_hierarchy(path) + # Start importing blend file. All objects that should be imported need to be copied from "data_from" to "data_to" with bpy.data.libraries.load(path, link=link) as (data_from, data_to): for data_block in data_blocks: @@ -57,6 +123,9 @@ def load_blend(path: str, obj_types: Optional[Union[List[str], str]] = None, nam else: raise Exception("No such data block: " + data_block) + if store_hierarchy: + build_collection_hierarchy(collection_hierarchy) + # Go over all imported objects again loaded_objects: List[Entity] = [] for data_block in data_blocks: @@ -65,9 +134,11 @@ def load_blend(path: str, obj_types: Optional[Union[List[str], str]] = None, nam for obj in getattr(data_to, data_block): # Check that the object type is desired if obj.type.lower() in obj_types: - # Link objects to the scene - bpy.context.collection.objects.link(obj) loaded_objects.append(convert_to_entity_subclass(obj)) + if store_hierarchy: + add_obj_to_hierarchy(obj, object_collections_hierarchy) + else: + bpy.context.collection.objects.link(obj) # If a camera was imported if obj.type == 'CAMERA':