RCG: new assets management

rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py

@@ -42,7 +42,7 @@ from bpy.props import (
 
 from .io_scene_json import export_json
 from .io_anim_ros2bag import set_animation_data
-from .io_entity_manager import switch_3d_entities
+from .io_assets_manager import add_to_asset_manager, add_links_assets
 
 bl_info = {
     'name': 'Robossembler Tools',
@@ -70,12 +70,12 @@ class addon_Properties(PropertyGroup):
                ]
     )
 
-    entity: EnumProperty(
+    assets_type: EnumProperty(
-        name='Entity',
+        name='Assets Type',
-        description='Selest 3d Entity',
+        description='Selest Assets Type',
-        items=[('hp', 'Highpoly', ''),
+        items=[('RENDER', 'render', ''),
-               ('mp', 'Modpoly', ''),
+               ('VISUAL', 'visual', ''),
-               ('lp', 'Lowpoly', '')
+               ('COLLISION', 'collision', '')
                ]
     )
 
@@ -95,10 +95,10 @@ class addon_Properties(PropertyGroup):
         subtype='FILE_PATH'
     )
 
-    refs_file_path: StringProperty(
+    store_dir: StringProperty(
-        name='Dir Path',
+        name='Nix Store Dir',
-        description='References library file path',
+        description='Resources root dir',
-        default='',
+        default='/media/disk/robossembler/project/pipeline/resources/',
         maxlen=1023,
         subtype='DIR_PATH'
     )
@@ -151,8 +151,8 @@ class RobossemblerPanel2(Panel):
 
 class RobossemblerPanel3(Panel):
     ''' Robossembler UI'''
-    bl_idname = 'ROBOSSEMBLER_PT_ENTITY_MANAGER'
+    bl_idname = 'ROBOSSEMBLER_PT_ASSETS_MANAGER'
-    bl_label = 'Switch 3d Entities'
+    bl_label = 'Manage assets types'
     bl_space_type = 'VIEW_3D'
     bl_region_type = 'UI'
     bl_category = 'Robossembler'
@@ -161,15 +161,15 @@ class RobossemblerPanel3(Panel):
         prop = context.scene.robossembler_properties
 
         layout = self.layout
-        layout.prop(prop, 'refs_file_path')
+        layout.prop(prop, 'store_dir')
-        layout.prop(prop, 'entity')
+        layout.prop(prop, 'assets_type')
 
         col = layout.column()
         col.alert = True
         col.scale_y = 2.0
-        col.operator('scene.manage_entities',
+        col.operator('scene.add_links_assets',
                      icon='ASSET_MANAGER',
-                     text='Switch Entities')
+                     text='Manage Assets')
 
 
 class RobossemblerOperator1(Operator):
@@ -207,18 +207,27 @@ class RobossemblerOperator2(Operator):
 
 class RobossemblerOperator3(Operator):
     '''Tooltip'''
-    bl_idname = 'scene.manage_entities'
+    bl_idname = 'scene.add_links_assets'
     bl_label = ''
-    bl_description = 'Switch visual 3d entities.'
+    bl_description = 'Add assets by type.'
     bl_options = {'REGISTER', 'UNDO'}
 
     def execute(self, context):
         prop = context.scene.robossembler_properties
 
-        file_path = os.path.realpath(bpy.path.abspath((prop.refs_file_path)))
+        assets_type = prop.assets_type
-        print('1'*10, file_path)
+
-        entity = prop.entity
+        # TODO
-        switch_3d_entities(context, file_path)
+        store_dir = os.path.realpath(bpy.path.abspath((prop.store_dir)))
+        resource_dirs = [
+            os.path.join(store_dir, name)
+            for name in os.listdir(store_dir)
+            if os.path.isdir(os.path.join(store_dir, name))
+            if 'assets.json' in os.listdir(os.path.join(store_dir, name))
+        ]
+
+        add_to_asset_manager(context, resource_dirs)
+        add_links_assets(context, assets_type)
 
         return {'FINISHED'}
 

rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py

@@ -0,0 +1,164 @@
+# coding: utf-8
+'''
+Copyright (C) 2024 Kurochkin Ilia @brothermechanic
+The Original Code by: Alexander Shushpanov @shalenikol
+Created by brothermechanic
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+'''
+
+import collections
+import logging
+import os
+from pathlib import Path
+from rosbags.highlevel import AnyReader
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def append_frame(idx: int, frames, links, frame_data) -> int:
+    ''' Append datalist for frame and increase frame index. '''
+    frame_data['links'] = links
+    frames.append(frame_data)
+    return idx + 1
+
+
+def gimbal_locks_filter(ros2bag_func):
+    ''' Finding and fix gimbal locks. '''
+    def wrapper(*args, **kwargs):
+        # finding gimbal locks
+        ros2bag_data = ros2bag_func(*args, **kwargs)
+        gimbal_locks = {}
+        for link in ros2bag_data['scene_links']:
+            gimbal_data = collections.defaultdict(list)
+            link_rotation = []
+            for frame_data in ros2bag_data['frames_data']:
+                link_rotation_pre = []
+                for link_data in frame_data['links']:
+                    if link_data['link'] == link:
+                        link_rotation_pre = link_rotation
+                        link_rotation = link_data['rot_wxyz']
+                if link_rotation_pre:
+                    idx = 0
+                    for axis_pre, axis_cur in zip(link_rotation_pre, link_rotation):
+                        if (round(axis_cur, 2) > 0 > round(axis_pre, 2)
+                                or round(axis_cur, 2) < 0 < round(axis_pre, 2)):
+                            gimbal_data[frame_data['id']].append(idx)
+                            gimbal_locks[link] = gimbal_data
+                            logger.info(
+                                'Gimbal lock detected for %s link %s axis at %s frame',
+                                link, idx, frame_data['id'])
+                        idx += 1
+
+        # fix gimbal locks
+        for link in gimbal_locks:
+            gimbal_frame = list(gimbal_locks[link])[0]
+            if len(gimbal_locks[link]) == 2:
+                gimbal_range = list(gimbal_locks[link])
+            else:
+                at_start = abs(ros2bag_data['frame_start'] - gimbal_frame)
+                at_end = abs(ros2bag_data['frame_end'] - gimbal_frame)
+                if at_start < at_end:
+                    gimbal_range = [ros2bag_data['frame_start'], gimbal_frame]
+                else:
+                    gimbal_range = [gimbal_frame, ros2bag_data['frame_end']]
+            for frame_data in ros2bag_data['frames_data'][gimbal_range[0]:gimbal_range[1]]:
+                for link_data in frame_data['links']:
+                    if link_data['link'] == link:
+                        for axis in gimbal_locks[link][gimbal_frame]:
+                            link_data['rot_wxyz'][axis] = (
+                                -1 * link_data['rot_wxyz'][axis])
+                            logger.debug(
+                                'Gimbal lock fixed for %s link %s axis %s',
+                                link, axis, frame_data['id'])
+            logger.info('Gimbal lock fixed for %s link', link)
+
+        return ros2bag_data
+
+    return wrapper
+
+
+#@gimbal_locks_filter
+def get_ros2bag_data(ros2bag_path: str,
+                       frame_start=0, frame_end=0, fps=30, target_topic='/tf') -> dict:
+    ''' Get animation data from Ros2Bag database '''
+
+    assert ros2bag_path.endswith('.db3'), (
+        'Please, check Ros2Bag file format and extension!')
+
+    scene_links = set()
+    frames_data = []
+    ros2bag_dir = os.path.split(ros2bag_path)[0]
+
+    with AnyReader([Path(ros2bag_dir.replace('\\', '/'))]) as ros2bag:
+        targets = [
+            x for x in ros2bag.connections if x.topic == target_topic]
+        # TODO Switch to timestamp instead of frame index
+        idx = 0
+        idx_pre = - 1
+        for connection, timestamp, rawdata in ros2bag.messages(connections=targets):
+            if idx != idx_pre:
+                frame_data = {'id': idx, 'timestamp': timestamp}
+                key_link = []
+                links = []
+                idx_pre = idx
+
+            tf_msg = ros2bag.deserialize(rawdata, connection.msgtype)
+            for t in tf_msg.transforms:
+                c_key = t.header.frame_id + t.child_frame_id
+                if key_link.count(c_key) > 0:
+                    idx = append_frame(idx, frames_data, links, frame_data)
+                else:
+                    key_link.append(c_key)
+                    scene_links.add(t.child_frame_id)
+                    if t.header.frame_id != 'world':
+                        scene_links.add(t.header.frame_id)
+                    link_data = {}
+                    link_data['parent'] = t.header.frame_id
+                    link_data['link'] = t.child_frame_id
+                    link_data['loc_xyz'] = [
+                        t.transform.translation.x,
+                        t.transform.translation.y,
+                        t.transform.translation.z]
+                    link_data['rot_wxyz'] = [
+                        t.transform.rotation.w,
+                        t.transform.rotation.x,
+                        t.transform.rotation.y,
+                        t.transform.rotation.z]
+                    links.append(link_data)
+
+            if frame_end and not idx < frame_end + 1:
+                break
+
+    assert len(frames_data) - 1 == frames_data[-1]['id'], (
+        'Ros2Bag database has dublicated frames!')
+
+    frame_end = frames_data[-1]['id']
+
+    return {'path': ros2bag_path,
+            'frame_start': frame_start,
+            'frame_end': frame_end,
+            'fps': fps,
+            'scene_links': list(scene_links),
+            'frames_data': frames_data}
+
+
+#if __name__ == '__main__':
+import json
+ros2bag_path = '/media/disk/robossembler/project/pipeline/projects/subset_0.db3'
+json_path = os.path.splitext(ros2bag_path)[0] + '.json'
+with open(json_path, "w", encoding='utf-8') as data_file:
+    json.dump(get_ros2bag_data(ros2bag_path), data_file, indent=4)
+logger.info('Database saved successfully to %s!', json_path)

rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py

@@ -18,7 +18,6 @@ Created by brothermechanic
     along with this program.  If not, see <http://www.gnu.org/licenses/>.
 '''
 
-import collections
 import logging
 import os
 from pathlib import Path
@@ -35,63 +34,7 @@ def append_frame(idx: int, frames, links, frame_data) -> int:
     return idx + 1
 
 
-def gimbal_locks_filter(ros2bag_func):
+def get_ros2bag_data(ros2bag_path: str,
-    ''' Finding and fix gimbal locks. '''
-    def wrapper(*args, **kwargs):
-        # finding gimbal locks
-        ros2bag_data = ros2bag_func(*args, **kwargs)
-        gimbal_locks = {}
-        for link in ros2bag_data['scene_links']:
-            gimbal_data = collections.defaultdict(list)
-            link_rotation = []
-            for frame_data in ros2bag_data['frames_data']:
-                link_rotation_pre = []
-                for link_data in frame_data['links']:
-                    if link_data['link'] == link:
-                        link_rotation_pre = link_rotation
-                        link_rotation = link_data['rot_wxyz']
-                if link_rotation_pre:
-                    idx = 0
-                    for axis_pre, axis_cur in zip(link_rotation_pre, link_rotation):
-                        if (round(axis_cur, 2) > 0 > round(axis_pre, 2)
-                                or round(axis_cur, 2) < 0 < round(axis_pre, 2)):
-                            gimbal_data[frame_data['id']].append(idx)
-                            gimbal_locks[link] = gimbal_data
-                            logger.info(
-                                'Gimbal lock detected for %s link %s axis at %s frame',
-                                link, idx, frame_data['id'])
-                        idx += 1
-
-        # fix gimbal locks
-        for link in gimbal_locks:
-            gimbal_frame = list(gimbal_locks[link])[0]
-            if len(gimbal_locks[link]) == 2:
-                gimbal_range = list(gimbal_locks[link])
-            else:
-                at_start = abs(ros2bag_data['frame_start'] - gimbal_frame)
-                at_end = abs(ros2bag_data['frame_end'] - gimbal_frame)
-                if at_start < at_end:
-                    gimbal_range = [ros2bag_data['frame_start'], gimbal_frame]
-                else:
-                    gimbal_range = [gimbal_frame, ros2bag_data['frame_end']]
-            for frame_data in ros2bag_data['frames_data'][gimbal_range[0]:gimbal_range[1]]:
-                for link_data in frame_data['links']:
-                    if link_data['link'] == link:
-                        for axis in gimbal_locks[link][gimbal_frame]:
-                            link_data['rot_wxyz'][axis] = (
-                                -1 * link_data['rot_wxyz'][axis])
-                            logger.debug(
-                                'Gimbal lock fixed for %s link %s axis %s',
-                                link, axis, frame_data['id'])
-            logger.info('Gimbal lock fixed for %s link', link)
-
-        return ros2bag_data
-
-    return wrapper
-
-
-@gimbal_locks_filter
-def get_animation_data(ros2bag_path: str,
                        frame_start=0, frame_end=0, fps=30, target_topic='/tf') -> dict:
     ''' Get animation data from Ros2Bag database '''
 
@@ -155,9 +98,10 @@ def get_animation_data(ros2bag_path: str,
             'frames_data': frames_data}
 
 
-if __name__ == '__main__':
+#if __name__ == '__main__':
 import json
-    in_file = '/media/disk/robossembler/project/collab/138-rosbag-to-blender/cg/blender/scripts/addons/rosbag-importer/rosbag2/rosbag2_2024_02_15-18_15_42_0.db3'
+ros2bag_path = '/media/disk/robossembler/project/pipeline/projects/subset_0.db3'
-    out_file = '/media/disk/robossembler/project/collab/138-rosbag-to-blender/cg/blender/scripts/addons/rosbag-importer/rosbag2/rosbag2_2024_02_15-18_15_42_0.json'
+json_path = os.path.splitext(ros2bag_path)[0] + '.json'
-    with open(out_file, "w", encoding='utf-8') as fh:
+with open(json_path, "w", encoding='utf-8') as data_file:
-        json.dump(get_animation_data(in_file), fh, indent=4)
+    json.dump(get_ros2bag_data(ros2bag_path), data_file, indent=4)
+logger.info('Database saved successfully to %s!', json_path)

rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_assets_manager/__init__.py

@@ -0,0 +1,166 @@
+#  ***** BEGIN GPL LICENSE BLOCK *****
+#
+#   Copyright (C) 2021-2024 Robossembler LLC
+#
+#   Created by Ilia Kurochkin (brothermechanic)
+#   contact:    brothermechanic@yandex.com
+#
+#   This file is part of Robossembler Framework
+#   project repo:    https://gitlab.com/robossembler/framework
+#
+#   Robossembler Framework is free software;
+#   you can redistribute it and/or modify
+#   it under the terms of the GNU General Public License
+#   as published by the Free Software Foundation; either version 3
+#   of the License, or (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU General Public License for more details.
+#
+#   You should have received a copy of the GNU General Public License
+#   along with this program; if not, see <https://www.gnu.org/licenses/>.
+#
+#  ***** END GPL LICENSE BLOCK *****
+#
+# coding: utf-8
+
+import logging
+import json
+import os
+
+import bpy
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+# TODO
+debug_table = {
+    'base_link': 'asm_start',
+    'ee_link': 'AsmTailLink_221214',
+    'main0_link': 'asmLinkMain221213',
+    'main1_link': 'asmLinkMain221213',
+    'fork0_link': 'asmMainFork_221220',
+    'fork1_link': 'asmMainFork_221220',
+    'fork2_link': 'asmMainFork_221220'
+}
+
+
+def recursive_layer_collection(layer_coll, coll_name):
+    found = None
+    if layer_coll.name == coll_name:
+        return layer_coll
+    for layer in layer_coll.children:
+        found = recursive_layer_collection(layer, coll_name)
+        if found:
+            return found
+    return False
+
+
+def add_to_asset_manager(context, resource_project_dirs):
+    asset_libraries = context.preferences.filepaths.asset_libraries
+    # allready added resources
+    stored_dirs = {alib.name: alib.path for alib in asset_libraries}
+
+    for resource_project_dir in resource_project_dirs:
+        # get database assets_data
+        assets_data_path = os.path.join(resource_project_dir, 'assets.json')
+        with open(assets_data_path, encoding='utf-8') as data:
+            assets_data = json.load(data)
+        # {'resource_RENDER': '/path/to/resource_project/blend'}
+        new_dirs = {
+            f'{os.path.basename(resource_project_dir)}_{asset_data["type"]}': (
+                os.path.join(resource_project_dir, os.path.dirname(asset_data['path'])))
+            for asset_data in assets_data
+        }
+        for new_dir_name in new_dirs:
+            if new_dir_name in stored_dirs:
+                if new_dirs[new_dir_name] == stored_dirs[new_dir_name]:
+                    continue
+                logger.warning(
+                    'Resource %s has unexpected path %s! Will be overwtitten!',
+                    new_dir_name, stored_dirs[new_dir_name])
+            new_dir = asset_libraries.new()
+            new_dir.path = new_dirs[new_dir_name]
+            new_dir.name = new_dir_name
+            new_dir.import_method = 'LINK'
+            logger.info('Added %s %s!', new_dir_name, new_dirs[new_dir_name])
+    #bpy.ops.wm.save_userpref()
+    return [os.path.basename(dir) for dir in resource_project_dirs]
+
+
+def add_links_assets(context, assets_type):
+    ''' '''
+    # check for main link collection
+    if not bpy.data.collections.get('link'):
+        raise Exception('No urdf/sdf assets in scene!')
+    # hide collections
+    for collection_name in ['render', 'visual', 'collision']:
+        if collection_name == assets_type.lower():
+            continue
+        if bpy.data.collections.get(collection_name):
+            bpy.data.collections[collection_name].hide_render = True
+            bpy.data.collections[collection_name].hide_viewport = True
+
+    if not bpy.data.collections.get(assets_type.lower()):
+        # create target collection
+        target_collection = bpy.data.collections.new(assets_type.lower())
+        context.scene.collection.children.link(target_collection)
+    else:
+        # cleanup target collection
+        target_collection = bpy.data.collections[assets_type.lower()]
+        list(map(bpy.data.objects.remove, target_collection.objects))
+        bpy.data.collections.remove(target_collection)
+        target_collection = bpy.data.collections.new(assets_type.lower())
+        context.scene.collection.children.link(target_collection)
+
+    active_collection = recursive_layer_collection(
+        context.view_layer.layer_collection,
+        target_collection.name)
+    context.view_layer.active_layer_collection = active_collection
+    # select assets
+    asset_libraries = context.preferences.filepaths.asset_libraries
+    actual_libraries = filter((lambda lib: os.path.isdir(lib.path)), asset_libraries)
+    selected_libraries = filter((lambda lib: (assets_type in lib.name)), actual_libraries)
+    # add assets
+    asset_paths = [
+        os.path.join(lib.path, blend)
+        for lib in selected_libraries
+        for blend in os.listdir(lib.path)
+        if os.path.isfile(os.path.join(lib.path, blend))
+    ]
+    link_arms = bpy.data.collections['link'].objects
+    for asset_path in asset_paths:
+        with bpy.data.libraries.load(asset_path, assets_only=True) as (
+            data_from, data_to):
+            asset_name, = data_from.collections
+        for link_arm in link_arms:
+            ###if asset_name != link_arm.name:
+            ###    continue
+            if debug_table.get(link_arm.name) != asset_name:
+                continue
+            ###
+            bpy.ops.object.select_all(action='DESELECT')
+            bpy.ops.wm.link(
+                filepath=asset_path,
+                directory=os.path.join(asset_path, 'Collection'),
+                filename=asset_name,
+                relative_path=True,
+                do_reuse_local_id=True,
+                active_collection=True,
+                autoselect=True
+            )
+            #asset_obj = context.object
+            asset_obj = context.selected_objects[0]
+            ###
+            import math
+            asset_obj.rotation_euler[2] = math.radians(90)
+            ###
+            assert asset_obj.instance_type == 'COLLECTION', 'Wrong Linking!'
+            asset_obj.parent = link_arm
+            asset_obj.parent_type = 'BONE'
+            asset_obj.parent_bone = link_arm.data.bones[0].name
+            logger.info('Added %s library!', asset_obj.name)
+    return True

rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py

@@ -1,88 +0,0 @@
-# coding: utf-8
-'''
-Copyright (C) 2024 brothermechanic@yandex.com
-
-Created by brothermechanic
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <http://www.gnu.org/licenses/>.
-'''
-
-import logging
-import os
-import bpy
-
-logger = logging.getLogger(__name__)
-
-
-def unlink_from_collections(context, obj):
-    ''' Unlinking object from all collections. '''
-    for col in list(bpy.data.collections) + [context.scene.collection]:
-        if obj.name in col.objects:
-            col.objects.unlink(obj)
-    return obj
-
-
-def link_library(lib_path, lib_type, lib_name):
-    '''  '''
-    bpy.ops.wm.link(
-        filepath=lib_path,
-        directory=os.path.join(lib_path, lib_type),
-        filename=lib_name,
-        relative_path=True,
-        do_reuse_local_id=True,
-        autoselect=True,
-        instance_collections=True,
-        instance_object_data=True
-        )
-    return bpy.data.objects[lib_name]
-
-
-def switch_3d_entities(context, lib_dir):
-    '''  '''
-    assert bpy.data.collections.get('visual'), 'No visual collection!'
-    entities_from = bpy.data.collections['visual'].objects
-    assert os.path.isdir(lib_dir), 'No libs dir {}!'.format(lib_dir)
-    lib_files = os.listdir(lib_dir)
-
-    hp_col = bpy.data.collections.new('Parts')
-    context.scene.collection.children.link(hp_col)
-    for entity_from in entities_from:
-        for lib_file in lib_files:
-            if '{}_hp.blend'.format(entity_from.name) != lib_file:
-                continue
-
-            entity_to = link_library(
-                lib_path=os.path.join(lib_dir, lib_file),
-                lib_type='Collection',
-                lib_name='{}_hp'.format(entity_from.name))
-
-            entity_to.empty_display_type = 'ARROWS'
-            entity_to.empty_display_size = 0.5
-
-            unlink_from_collections(context, entity_to)
-            hp_col.objects.link(entity_to)
-
-            entity_to.location = entity_from.location
-            entity_from.rotation_mode = entity_to.rotation_mode = 'QUATERNION'
-            entity_to.rotation_quaternion = entity_from.rotation_quaternion
-            entity_to.scale = entity_from.scale
-
-            entity_to.parent = entity_from.parent
-
-            logger.info('Entity %s changed to %s!', entity_from.name, entity_to.name)
-
-    bpy.data.collections['visual'].hide_render = True
-    bpy.data.collections['visual'].hide_viewport = True
-
-    return True