From ba6a8e9d082ee471acc4ef20fa448a607aee8982 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Thu, 25 Apr 2024 10:24:17 +0300 Subject: [PATCH] RCG: new assets management --- .../scripts/addons/Robossembler/__init__.py | 55 +++--- .../io_anim_ros2bag/_ros2bag_parser.py | 164 +++++++++++++++++ .../io_anim_ros2bag/ros2bag_parser.py | 72 +------- .../io_assets_manager/__init__.py | 166 ++++++++++++++++++ .../io_entity_manager/__init__.py | 88 ---------- 5 files changed, 370 insertions(+), 175 deletions(-) create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_assets_manager/__init__.py delete mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py index 1c212fa..80c7803 100644 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py +++ b/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( - name='Entity', - description='Selest 3d Entity', - items=[('hp', 'Highpoly', ''), - ('mp', 'Modpoly', ''), - ('lp', 'Lowpoly', '') + assets_type: EnumProperty( + name='Assets Type', + description='Selest Assets Type', + items=[('RENDER', 'render', ''), + ('VISUAL', 'visual', ''), + ('COLLISION', 'collision', '') ] ) @@ -95,10 +95,10 @@ class addon_Properties(PropertyGroup): subtype='FILE_PATH' ) - refs_file_path: StringProperty( - name='Dir Path', - description='References library file path', - default='', + store_dir: StringProperty( + name='Nix Store Dir', + description='Resources root dir', + 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_label = 'Switch 3d Entities' + bl_idname = 'ROBOSSEMBLER_PT_ASSETS_MANAGER' + 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, 'entity') + layout.prop(prop, 'store_dir') + 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))) - print('1'*10, file_path) - entity = prop.entity - switch_3d_entities(context, file_path) + assets_type = prop.assets_type + + # TODO + 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'} diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py new file mode 100644 index 0000000..cfe239c --- /dev/null +++ b/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 . +''' + +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) diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py index 7908d3b..3fa8ccc 100644 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py +++ b/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 . ''' -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): - ''' 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, +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 ''' @@ -155,9 +98,10 @@ def get_animation_data(ros2bag_path: str, 'frames_data': frames_data} -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' - 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' - with open(out_file, "w", encoding='utf-8') as fh: - json.dump(get_animation_data(in_file), fh, indent=4) +#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) diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_assets_manager/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_assets_manager/__init__.py new file mode 100644 index 0000000..87a994c --- /dev/null +++ b/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 . +# +# ***** 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 diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py deleted file mode 100644 index 00cce18..0000000 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py +++ /dev/null @@ -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 . -''' - -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