From ec2d5cae0a2a5033b690b8a4c35a95bba03eb5f7 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Wed, 17 Apr 2024 13:19:58 +0300 Subject: [PATCH 1/3] Add Robossembler Addon from DEPRECATED CG overlay --- .../scripts/addons/Robossembler/__init__.py | 249 ++++++++++++++++++ .../Robossembler/io_anim_ros2bag/__init__.py | 66 +++++ .../io_anim_ros2bag/ros2bag_parser.py | 163 ++++++++++++ .../io_entity_manager/__init__.py | 88 +++++++ .../Robossembler/io_scene_json/__init__.py | 189 +++++++++++++ .../Robossembler/io_scene_json/model_md5.py | 30 +++ .../Robossembler/io_scene_json/model_name.py | 18 ++ .../Robossembler/io_scene_json/model_paths.py | 12 + .../scripts/startup/cg_environment.py | 42 +++ 9 files changed, 857 insertions(+) create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py 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_entity_manager/__init__.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/__init__.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_md5.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_name.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_paths.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/startup/cg_environment.py diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py new file mode 100644 index 0000000..1c212fa --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py @@ -0,0 +1,249 @@ +# ***** BEGIN GPL LICENSE BLOCK ***** +# +# 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 +# and write to the Free Software Foundation, Inc., 51 Franklin Street, +# Fifth Floor, Boston, MA 02110-1301, USA.. +# +# The Original Code is Copyright (C) 2023 by Kurochkin Ilia ### +# All rights reserved. +# +# Contact: brothermechanic@yandex.com ### +# Information: https://gitlab.com/robossembler ### +# +# The Original Code is: all of this file. +# +# ***** END GPL LICENSE BLOCK ***** +# +# -*- coding: utf-8 -*- + +import os +import bpy + +from bpy.types import ( + Panel, + Operator, + PropertyGroup) + +from bpy.props import ( + StringProperty, + EnumProperty, + PointerProperty) + +from .io_scene_json import export_json +from .io_anim_ros2bag import set_animation_data +from .io_entity_manager import switch_3d_entities + +bl_info = { + 'name': 'Robossembler Tools', + 'author': 'brothermechanic@gmail.com', + 'version': (0, 2), + 'blender': (4, 2, 0), + 'location': '3D View > Toolbox', + 'description': 'Robossembler pipeline tools', + 'warning': '', + 'wiki_url': '', + 'tracker_url': 'https://gitlab.com/robossembler', + 'category': 'Robossembler', +} + + +class addon_Properties(PropertyGroup): + + engine: EnumProperty( + name='Engine', + description='Selest Target Engine', + items=[('BULLET', 'Bullet', ''), + ('ODE', 'O D E', ''), + ('SIMBODY', 'Simbody', ''), + ('OPENSIM', 'OpenSim', '') + ] + ) + + entity: EnumProperty( + name='Entity', + description='Selest 3d Entity', + items=[('hp', 'Highpoly', ''), + ('mp', 'Modpoly', ''), + ('lp', 'Lowpoly', '') + ] + ) + + json_file_path: StringProperty( + name='File Path', + description='Input/output Json file path', + default='', + maxlen=1023, + subtype='FILE_PATH' + ) + + ros2bag_file_path: StringProperty( + name='File Path', + description='Input Ros2Bag file path', + default='', + maxlen=1023, + subtype='FILE_PATH' + ) + + refs_file_path: StringProperty( + name='Dir Path', + description='References library file path', + default='', + maxlen=1023, + subtype='DIR_PATH' + ) + + +class RobossemblerPanel1(Panel): + ''' Robossembler UI''' + bl_idname = 'ROBOSSEMBLER_PT_EXPORT_JSON' + bl_label = 'Export Scene as Json' + bl_space_type = 'VIEW_3D' + bl_region_type = 'UI' + bl_category = 'Robossembler' + + def draw(self, context): + prop = context.scene.robossembler_properties + + layout = self.layout + layout.prop(prop, 'json_file_path') + layout.prop(prop, 'engine') + + col = layout.column() + col.alert = True + col.scale_y = 2.0 + col.operator('scene.export_json', + icon='WORLD_DATA', + text='Export Scene') + + +class RobossemblerPanel2(Panel): + '''Doc''' + bl_idname = 'ROBOSSEMBLER_PT_IMPORT_ANIMATION' + bl_label = 'Import Ros2Bag Animation' + bl_space_type = 'VIEW_3D' + bl_region_type = 'UI' + bl_category = 'Robossembler' + + def draw(self, context): + prop = context.scene.robossembler_properties + + layout = self.layout + layout.prop(prop, 'ros2bag_file_path') + + col = layout.column() + col.alert = True + col.scale_y = 2.0 + col.operator('scene.import_ros2bag', + icon='ACTION', + text='Import Animation') + + +class RobossemblerPanel3(Panel): + ''' Robossembler UI''' + bl_idname = 'ROBOSSEMBLER_PT_ENTITY_MANAGER' + bl_label = 'Switch 3d Entities' + bl_space_type = 'VIEW_3D' + bl_region_type = 'UI' + bl_category = 'Robossembler' + + def draw(self, context): + prop = context.scene.robossembler_properties + + layout = self.layout + layout.prop(prop, 'refs_file_path') + layout.prop(prop, 'entity') + + col = layout.column() + col.alert = True + col.scale_y = 2.0 + col.operator('scene.manage_entities', + icon='ASSET_MANAGER', + text='Switch Entities') + + +class RobossemblerOperator1(Operator): + '''Tooltip''' + bl_idname = 'scene.export_json' + bl_label = '' + bl_description = 'Export scene liks to json config operator.' + bl_options = {'REGISTER', 'UNDO'} + + def execute(self, context): + prop = context.scene.robossembler_properties + + file_path = os.path.realpath(bpy.path.abspath((prop.json_file_path))) + physics_engine = prop.engine + export_json(context, file_path, physics_engine) + + return {'FINISHED'} + + +class RobossemblerOperator2(Operator): + '''Tooltip''' + bl_idname = 'scene.import_ros2bag' + bl_label = '' + bl_description = 'Import Ros2Bag animation to scene and apply it to liks.' + bl_options = {'REGISTER', 'UNDO'} + + def execute(self, context): + prop = context.scene.robossembler_properties + + file_path = os.path.realpath(bpy.path.abspath((prop.ros2bag_file_path))) + set_animation_data(context, file_path) + + return {'FINISHED'} + + +class RobossemblerOperator3(Operator): + '''Tooltip''' + bl_idname = 'scene.manage_entities' + bl_label = '' + bl_description = 'Switch visual 3d entities.' + 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) + + return {'FINISHED'} + + +classes = ( + RobossemblerPanel1, + RobossemblerPanel2, + RobossemblerPanel3, + RobossemblerOperator1, + RobossemblerOperator2, + RobossemblerOperator3, + addon_Properties) + + +def register(): + for cls in classes: + bpy.utils.register_class(cls) + bpy.types.Scene.robossembler_properties = PointerProperty(type=addon_Properties) + + +def unregister(): + for cls in classes: + bpy.utils.unregister_class(cls) + del bpy.types.Scene.robossembler_properties + + +if __name__ == '__main__': + register() diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py new file mode 100644 index 0000000..e94b048 --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py @@ -0,0 +1,66 @@ +# 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 json +import math +import bpy + +from .ros2bag_parser import get_animation_data + +logger = logging.getLogger(__name__) + + +def set_animation_data(context, ros2bag_path, frame_start=0, frame_end=0, fps=30): + ''' Set animation data from Ros2Bag database ''' + scene = context.scene + + ros2bag_data = get_animation_data( + ros2bag_path=ros2bag_path, frame_start=frame_start, frame_end=frame_end, fps=fps) + + scene.frame_start = ros2bag_data['frame_start'] + scene.frame_end = ros2bag_data['frame_end'] + scene.render.fps = ros2bag_data['fps'] + lost_links = [item for item in ros2bag_data['scene_links'] if not bpy.data.objects.get(item)] + if lost_links: + logger.warning('Link(s) not found in current scene: %s', lost_links) + + for frame_data in ros2bag_data['frames_data']: + frame_data['id'] + for link_data in frame_data['links']: + if link_data['link'] in lost_links: + continue + + parent_location = ((0, 0, 0) if link_data['parent'] == 'world' + else bpy.data.objects[link_data['parent']].location[:]) + + link = bpy.data.objects[link_data['link']] + link.rotation_mode = 'QUATERNION' + + link.location = [a + b for a, b in zip(link_data['loc_xyz'], parent_location)] + link.rotation_quaternion = link_data['rot_wxyz'] + + link.keyframe_insert(data_path='location', frame=frame_data['id']) + link.keyframe_insert(data_path='rotation_quaternion', frame=frame_data['id']) + + link.animation_data.action.name = '{}_action'.format(link_data['link']) + + logger.info('Setup of %s links is finished!', len(ros2bag_data['scene_links'])) + return True 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..7908d3b --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py @@ -0,0 +1,163 @@ +# 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_animation_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 + 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) 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 new file mode 100644 index 0000000..00cce18 --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_entity_manager/__init__.py @@ -0,0 +1,88 @@ +# 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 diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/__init__.py new file mode 100644 index 0000000..732f24b --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/__init__.py @@ -0,0 +1,189 @@ +# ***** BEGIN GPL LICENSE BLOCK ***** +# +# 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 +# and write to the Free Software Foundation, Inc., 51 Franklin Street, +# Fifth Floor, Boston, MA 02110-1301, USA.. +# +# The Original Code is Copyright (C) 2023 by Kurochkin Ilia ### +# All rights reserved. +# +# Contact: brothermechanic@yandex.com ### +# Information: https://gitlab.com/robossembler ### +# +# The Original Code is: all of this file. +# +# ***** END GPL LICENSE BLOCK ***** +# +# -*- coding: utf-8 -*- +''' +DESCRIPTION. +Collect all root objects in blender scene and export as json scene configuration. +''' +__version__ = '0.1' + +import collections +import json +import logging +import math +import os + +from .model_paths import get_urdf_sdf_model_paths +from .model_name import get_sdf_urdf_model_name +from .model_md5 import create_file_md5 + + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + +def export_json(context, scene_config_path, physics_engine): + ''' Export scene config to json. ''' + # open conf file + if os.path.isfile(scene_config_path): + with open(scene_config_path, 'r', encoding='utf-8') as conf_file: + world_dic = json.load(conf_file) + else: + world_dic = collections.defaultdict(list) + models_dir = os.path.join(os.path.dirname(scene_config_path), 'models') + model_paths = get_urdf_sdf_model_paths(models_dir) + for model_path in model_paths: + model_dic = {} + # 1 + model_dic['name'] = get_sdf_urdf_model_name(model_path) + # 2 + model_dic['id'] = create_file_md5(model_path) + # 3 + relative_model_path = model_path.split( + os.path.dirname(scene_config_path))[1][1:] + model_dic['path'] = relative_model_path + world_dic['models'].append(model_dic) + + # collect all non parented objs + roots = [obj for obj in context.scene.objects if not obj.parent] + + # collect instances + for instance in roots: + instance_dic = {} + # 1 + instance_dic['model_name'] = '' + if instance.get('model/name'): + instance_dic['model_name'] = instance['model/name'] + # 2 + instance_dic['model_id'] = '' + for item in world_dic['models']: + if item['name'] != instance_dic['model_name']: + continue + instance_dic['model_id'] = item['id'] + # 2 + instance_dic['id'] = '' + # 3 + instance_dic['pose'] = {} + instance_dic['pose']['x'] = round(instance.location[0], 3) + instance_dic['pose']['y'] = round(instance.location[1], 3) + instance_dic['pose']['z'] = round(instance.location[2], 3) + instance_dic['pose']['roll'] = round(math.degrees(instance.rotation_euler[0]), 3) + instance_dic['pose']['pitch'] = round(math.degrees(instance.rotation_euler[1]), 3) + instance_dic['pose']['yaw'] = round(math.degrees(instance.rotation_euler[2]), 3) + # 4 + if instance.scale[0] != sum(instance.scale[:]) / 3: + logger.warning('Asset instance {} is not uniformly scaled!'.format( + instance.name)) + instance_dic['scale'] = round(instance.scale[0], 3) + # 5 + # if assembly + if instance_dic['model_name']: + instance_dic['type'] = 'asset' + links = instance.children_recursive + links.insert(0, instance) + instance_dic['link_names'] = [] + instance_dic['links'] = {} + transforms = 0.0 + for link in links: + child = {} + if link.type != 'ARMATURE': + continue + bone = link.pose.bones[0] + rotation_mode = bone.rotation_mode + bone.rotation_mode = 'XYZ' + child['pose'] = {} + child['pose']['x'] = round(bone.location[0], 6) + child['pose']['y'] = round(bone.location[1], 6) + child['pose']['z'] = round(bone.location[2], 6) + child['pose']['roll'] = round(math.degrees(bone.rotation_euler[0]), 6) + child['pose']['pitch'] = round(math.degrees(bone.rotation_euler[1]), 6) + child['pose']['yaw'] = round(math.degrees(bone.rotation_euler[2]), 6) + bone.rotation_mode = rotation_mode + if bone.scale[0] != sum(bone.scale[:]) / 3: + logger.warning( + 'Link {} of asset {} is not uniformly scaled!'.format( + link.name, asset_dic['name'])) + child['scale'] = round(bone.scale[0], 3) + transforms += round(sum( + [child['pose']['x'], child['pose']['y'], child['pose']['z'], + child['pose']['roll'], child['pose']['pitch'], child['pose']['yaw']], + (child['scale'] - 1.0)), + 6) + if '.' in link.name: + link_name = link.name.rpartition('.')[0] + else: + link_name = link.name + + instance_dic['link_names'].append(link_name) + instance_dic['links'][link_name] = child + + if transforms == 0.0: + instance_dic.pop('link_names') + instance_dic.pop('links') + + # 6 + instance_dic['parent'] = 'world' + # if light + if instance.type == 'LIGHT': + instance_dic['type'] = instance.type.lower() + if instance.data.type == 'POINT': + instance_dic['light_type'] = 'point' + elif instance.data.type == 'SUN': + instance_dic['light_type'] = 'directional' + elif instance.data.type == 'SPOT': + instance_dic['light_type'] = 'spot' + instance_dic['spot_angle'] = round( + math.degrees(instance.data.spot_size), 1) + else: + logger.warning( + 'Unsupported light type {} on instance {}!'.format( + instance.data.type, instance.name)) + instance_dic['intencity'] = instance.data.energy + instance_dic['diffuse'] = list(instance.data.color) + # if camera + if instance.type == 'CAMERA': + instance_dic['type'] = instance.type.lower() + instance_dic['focal_length'] = instance.data.lens + instance_dic['sensor_size'] = [instance.data.sensor_width, + instance.data.sensor_height] + + world_dic['instances'].append(instance_dic) + + # select physics engine + world_dic['physics'] = {} + world_dic['physics']['engine_name'] = physics_engine + world_dic['physics']['gravity'] = {} + world_dic['physics']['gravity']['x'] = round(context.scene.gravity[0], 3) + world_dic['physics']['gravity']['y'] = round(context.scene.gravity[1], 3) + world_dic['physics']['gravity']['z'] = round(context.scene.gravity[2], 3) + + # write conf file + with open(scene_config_path, 'w', encoding='utf-8') as world_conf_file: + json.dump(world_dic, world_conf_file, ensure_ascii=False, indent=4) + logger.info('Scene configuration in json file was completed successfully!') + return scene_config_path diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_md5.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_md5.py new file mode 100644 index 0000000..e76d6fd --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_md5.py @@ -0,0 +1,30 @@ +# coding: utf-8 + +import hashlib +import os + + +def md5(file_path): + ''' Generate md5 hash. ''' + hash_md5 = hashlib.md5() + with open(file_path, "rb") as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + return hash_md5.hexdigest() + + +def create_file_md5(file_path): + ''' Get md5_file_hash of file and store to md5 file. ''' + md5_file_path = f'{file_path}.md5' + md5_file_hash = md5(file_path) + # create md5_file + if not os.path.isfile(md5_file_path): + with open(md5_file_path, 'w', encoding='utf-8') as md5_file: + md5_file.write(md5_file_hash) + else: + with open(md5_file_path, 'r', encoding='utf-8') as md5_file: + md5_file_stored = md5_file.read() + # check md5_file + assert md5_file_hash == md5_file_stored, ( + "Mosel's md5 don't matched with stored %s.md5 file", file_path) + return md5_file_hash diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_name.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_name.py new file mode 100644 index 0000000..4fb0a50 --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_name.py @@ -0,0 +1,18 @@ +# coding: utf-8 + +import logging +import xml.etree.ElementTree as ET + +logger = logging.getLogger(__name__) + + +def get_sdf_urdf_model_name(model_path): + mytree = ET.parse(model_path) + myroot = mytree.getroot() + for elem in myroot.iter(): + if elem.tag in ('link', 'joint', 'material', 'visual', 'collision'): + continue + if not elem.attrib.get('name'): + continue + return elem.attrib.get('name') + return logger.warning('Model %s do not have a name!', model_path) diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_paths.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_paths.py new file mode 100644 index 0000000..5f0c0fa --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_scene_json/model_paths.py @@ -0,0 +1,12 @@ +# coding: utf-8 + +import glob + + +def get_urdf_sdf_model_paths(models_dir): + ''' Collect models paths. ''' + model_paths = [] + for file_name in glob.glob(f'{models_dir}/**', recursive=True): + if file_name.endswith('.urdf') or file_name.endswith('.sdf'): + model_paths.append(file_name.replace('\\', '/')) + return model_paths diff --git a/rcg_pipeline/rcg_pipeline/scripts/startup/cg_environment.py b/rcg_pipeline/rcg_pipeline/scripts/startup/cg_environment.py new file mode 100644 index 0000000..0c69d21 --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/startup/cg_environment.py @@ -0,0 +1,42 @@ +#!/usr/bin/python3 +# -*- coding: utf-8 -*- +# Copyright (C) 2023 Ilia Kurochkin +# +# 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. + +__doc__ = 'Environment for use Blender with Robossembler Framework' +__version__ = '0.1' + +import os +import bpy +from bpy.app.handlers import persistent +import addon_utils + + +@persistent +def rs_env(): + '''Scripts environment''' + rs_blender_scripts_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) + if rs_blender_scripts_dir not in bpy.utils.script_paths_pref(): + script_directories = bpy.context.preferences.filepaths.script_directories + new_dir = script_directories.new() + new_dir.directory = rs_blender_scripts_dir + new_dir.name = 'RS_BLENDER_SCRIPTS_DIR' + bpy.ops.wm.save_userpref() + bpy.ops.wm.quit_blender() + else: + if not addon_utils.check('BakeWrangler')[0]: + addon_utils.enable('BakeWrangler', default_set=True) + bpy.ops.wm.save_userpref() + + print('Robossembler Framework Environment activated!') + +rs_env() From ba6a8e9d082ee471acc4ef20fa448a607aee8982 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Thu, 25 Apr 2024 10:24:17 +0300 Subject: [PATCH 2/3] 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 From bafa332cacc8a4fd7213a1ca8f67e1f82e0830b7 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Mon, 1 Jul 2024 14:12:33 +0300 Subject: [PATCH 3/3] robot builder wip --- rcg_pipeline/README.md | 2 +- rcg_pipeline/rcg_pipeline/__init__.py | 14 +- rcg_pipeline/rcg_pipeline/render_asset.py | 17 +- rcg_pipeline/rcg_pipeline/robot_asset.py | 127 ++++++++++++ .../scripts/addons/Robossembler/__init__.py | 2 +- .../addons/Robossembler/export_assembly.py | 71 +++++++ .../Robossembler/io_anim_ros2bag/__init__.py | 45 +++-- .../io_anim_ros2bag/_ros2bag_parser.py | 164 ---------------- .../io_anim_ros2bag/ros2bag_parser.py | 183 ++++++++++++------ .../rcg_pipeline/utils/collection_tools.py | 12 ++ 10 files changed, 381 insertions(+), 256 deletions(-) create mode 100644 rcg_pipeline/rcg_pipeline/robot_asset.py create mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/export_assembly.py delete mode 100644 rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py diff --git a/rcg_pipeline/README.md b/rcg_pipeline/README.md index e4791f5..8786740 100644 --- a/rcg_pipeline/README.md +++ b/rcg_pipeline/README.md @@ -31,5 +31,5 @@ import rcg_pipeline project_dir = '/path/to/' rcg_pipeline.libs.generate_libs_database(project_dir) -rcg_pipeline.render_asset.build_render_assets(project_dir) +rcg_pipeline.render_assets.build_render_assets(project_dir) ``` diff --git a/rcg_pipeline/rcg_pipeline/__init__.py b/rcg_pipeline/rcg_pipeline/__init__.py index 14bd09b..9d54d15 100644 --- a/rcg_pipeline/rcg_pipeline/__init__.py +++ b/rcg_pipeline/rcg_pipeline/__init__.py @@ -39,12 +39,14 @@ __email__ = 'brothermechanic@yandex.com' __copyright__ = 'Copyright (C) 2021-2024 Robossembler LLC' __url__ = ['https://robossembler.org'] __license__ = 'GPL-3' -#__all__ = ['libs', 'render_asset', 'rcg_full_pipeline'] +#__all__ = ['libs', 'render_assets', 'rcg_full_pipeline'] import logging from . import libs from . import render_asset +from . import robot_asset +#from . import world_assets logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -56,7 +58,15 @@ def rcg_full_pipeline(project_dir): # 1 generate libs libs.generate_libs_database(project_dir) - # 2 build render assets + # 2.1 build render assets render_asset.build_render_assets(project_dir) return True + + +# 3 build robot assets +def build_robot_assets(robot_dirs): + for robot_dir in robot_dirs: + robot_asset.build_robot_asset(robot_dir) +# 4 build world assets +# world_assets.build_world_assets(assembly_dirs, world_dirs) diff --git a/rcg_pipeline/rcg_pipeline/render_asset.py b/rcg_pipeline/rcg_pipeline/render_asset.py index d6954dc..db98fd1 100644 --- a/rcg_pipeline/rcg_pipeline/render_asset.py +++ b/rcg_pipeline/rcg_pipeline/render_asset.py @@ -60,9 +60,12 @@ def recursive_layer_collection(layer_coll, coll_name): return False -def assembly_builder(item, libs_data, libs_data_dir, collection=None, parent=None): +def assembly_builder(item, libs_data, collection=None, parent=None): ''' ''' + if not item.get('type'): + item['type'] = 'PART' + if not collection: collection = bpy.context.scene.collection @@ -99,11 +102,10 @@ def assembly_builder(item, libs_data, libs_data_dir, collection=None, parent=Non if not loc.library] if local_obj: local_obj[0].name += '_loc' - item_file_path = os.path.join(libs_data_dir, item_data[0]['path']) # TODO already linked bpy.ops.wm.link( - filepath=item_file_path, - directory=os.path.join(item_file_path, 'Collection'), + filepath=item_data[0]['path'], + directory=os.path.join(item_data[0]['path'], 'Collection'), filename=item['base_name'], relative_path=True, do_reuse_local_id=True, @@ -137,7 +139,7 @@ def assembly_builder(item, libs_data, libs_data_dir, collection=None, parent=Non if item.get('children'): for child_item in item.get('children'): assembly_builder( - child_item, libs_data, libs_data_dir, collection, parent=item_obj) + child_item, libs_data, collection, parent=item_obj) return True @@ -225,8 +227,11 @@ def build_render_assets(project_dir): render_collection.name) bpy.context.view_layer.active_layer_collection = active_collection + # solve libs paths + for lib_data in libs_data: + lib_data['path'] = os.path.join(project_dir, lib_data['path']) # build original hierarchy - assembly_builder(tree_item, libs_data, project_dir, render_collection) + assembly_builder(tree_item, libs_data, render_collection) # rebuild to LCS hierarchy assembly_rebuilder() diff --git a/rcg_pipeline/rcg_pipeline/robot_asset.py b/rcg_pipeline/rcg_pipeline/robot_asset.py new file mode 100644 index 0000000..9f877a4 --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/robot_asset.py @@ -0,0 +1,127 @@ +# ***** 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 +''' +DESCRIPTION. +Generate render asset from assembly tree and CG libs database. +''' + +__version__ = '1.0' + +import logging +import json +import os +import shutil + +import bpy + +from .utils.collection_tools import remove_collections_with_objects +from .utils.cleanup_orphan_data import cleanup_orphan_data +from .render_asset import (assembly_builder, recursive_layer_collection) + +logger = logging.getLogger(__name__) + + +def collect_asset_base_paths(tree_data): + yield tree_data['base_path'] + for asset_data in tree_data['children']: + yield from collect_asset_base_paths(asset_data) + +def solve_assets_path(assets_data, asset_base_path): + for asset in assets_data: + asset['path'] = os.path.join(asset_base_path, asset['path']) + return assets_data + + +def build_robot_asset(robot_dir): + ''' + ''' + robot_dir = os.path.normpath(robot_dir) + robot_name = os.path.basename(robot_dir) + robot_data_path = os.path.join(robot_dir, robot_name + '.json') + blend_path = os.path.join(robot_dir, robot_name + '.blend') + if not os.path.exists(robot_data_path): + raise Exception('No robot database found! Check %s directory' % robot_dir) + with open(robot_data_path, encoding='utf-8') as data: + robot_data = json.load(data) + + # robot dir cleanup + if os.path.exists(blend_path): + os.remove(blend_path) + if os.path.exists(blend_path + '1'): + os.remove(blend_path + '1') + + # start from stratch + bpy.ops.wm.read_homefile() + remove_collections_with_objects() + cleanup_orphan_data() + + # create robot collection + robot_collection = bpy.data.collections.new(robot_name) + bpy.context.scene.collection.children.link(robot_collection) + active_collection = recursive_layer_collection( + bpy.context.view_layer.layer_collection, + robot_collection.name) + bpy.context.view_layer.active_layer_collection = active_collection + + asset_base_paths = list(set(collect_asset_base_paths(robot_data))) + all_assets_data =[] + for asset_base_path in asset_base_paths: + assets_data_path = os.path.join(asset_base_path, 'assets.json') + if not os.path.exists(assets_data_path): + raise Exception('No assets database found! Check %s directory' % asset_base_path) + with open(assets_data_path, encoding='utf-8') as data: + assets_data = json.load(data) + # solve assets paths + for asset in assets_data: + asset['path'] = os.path.join(asset_base_path, asset['path']) + all_assets_data += assets_data + assembly_builder(robot_data, all_assets_data) + + # mark as asset + robot_collection.asset_mark() + # TODO collection thumbnail + + bpy.ops.wm.save_as_mainfile(filepath=blend_path) + + robot_data = [] + robot_data.append( + { + 'type': 'ROBOT', + 'name': robot_name, + 'path': os.path.relpath(blend_path, robot_dir), + 'thumbnail': '' + } + ) + logger.info('Robot Asset %s was generated!', robot_name) + # write db file + robot_data_path = os.path.join(robot_dir, 'robot.json') + with open(robot_data_path, 'w', encoding='utf-8') as assets_data_file: + json.dump(robot_data, assets_data_file, ensure_ascii=False, indent=4) + logger.info('Database saved successfully to %s!', robot_data_path) + + return blend_path diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py index 80c7803..f9e098b 100644 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/__init__.py @@ -47,7 +47,7 @@ from .io_assets_manager import add_to_asset_manager, add_links_assets bl_info = { 'name': 'Robossembler Tools', 'author': 'brothermechanic@gmail.com', - 'version': (0, 2), + 'version': (1, 0), 'blender': (4, 2, 0), 'location': '3D View > Toolbox', 'description': 'Robossembler pipeline tools', diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/export_assembly.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/export_assembly.py new file mode 100644 index 0000000..4bc14af --- /dev/null +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/export_assembly.py @@ -0,0 +1,71 @@ +import logging +import os +import json + +import bpy + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + +def round_locations(objs): + ''' Geting location of objects and round it. ''' + for obj in objs: + for idx, axis in enumerate(obj.location[:]): + obj.location[idx] = round(axis, 5) + return objs + + +def assembly_tree(obj, dict_tree) -> dict: + ''' ''' + # collect lib assets only + if not obj.instance_type == 'COLLECTION': + return False + # collect name + dict_tree['name'] = obj.name + dict_tree['base_name'] = obj.instance_collection.name + # dict_tree['file_name'] = obj.instance_collection.library.name + file_path = os.path.realpath(bpy.path.abspath( + obj.instance_collection.library.filepath)) + dict_tree['base_path'] = file_path.split('/assets')[0] + # collect transforms + obj.rotation_mode = 'QUATERNION' + dict_tree['pose'] = [ + {'loc_xyz': [round(axis, 5) for axis in obj.location]}, + {'rot_wxyz': [round(axis, 5) for axis in obj.rotation_quaternion]} + ] + # collect children + dict_tree['children'] = [] + for child in obj.children: + # skip hidden objects + if child.hide_get(): + continue + # skip nonlib objects + if child.instance_type == 'COLLECTION': + dict_tree['children'].append({}) + assembly_tree(child, dict_tree['children'][-1]) + return True + + +def export_assembly_config(): + ''' ''' + assert len(bpy.context.scene.collection.children) == 1, ( + 'Assembly should have only one collection!') + asm_collection = bpy.context.scene.collection.children[0] + root = [obj for obj in asm_collection.objects if not obj.parent] + assert len(root) == 1, 'Assembly should have only one root!' + dict_tree = {} + assembly_tree(root[0], dict_tree) + # write file + project_dir = '/media/disk/robossembler/project/pipeline/assemblies/arm/' + config_name = asm_collection.name + assembly_path = os.path.join(project_dir, config_name + '.json') + with open(assembly_path, 'w', encoding='utf-8') as json_file: + json.dump(dict_tree, json_file, ensure_ascii=False, indent=4) + logger.info('Assembly tree saved successfully to %s!', assembly_path) + + return assembly_path + + +round_locations(bpy.data.objects) +export_assembly_config() diff --git a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py index e94b048..08d15bb 100644 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py +++ b/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py @@ -28,39 +28,42 @@ from .ros2bag_parser import get_animation_data logger = logging.getLogger(__name__) -def set_animation_data(context, ros2bag_path, frame_start=0, frame_end=0, fps=30): +def set_animation_data(context, ros2bag_path): ''' Set animation data from Ros2Bag database ''' scene = context.scene - ros2bag_data = get_animation_data( - ros2bag_path=ros2bag_path, frame_start=frame_start, frame_end=frame_end, fps=fps) + ros2bag_data = get_animation_data(ros2bag_path=ros2bag_path) scene.frame_start = ros2bag_data['frame_start'] scene.frame_end = ros2bag_data['frame_end'] scene.render.fps = ros2bag_data['fps'] - lost_links = [item for item in ros2bag_data['scene_links'] if not bpy.data.objects.get(item)] - if lost_links: - logger.warning('Link(s) not found in current scene: %s', lost_links) + for topic in ros2bag_data['topics']: + # TODO if topic['name'] - for frame_data in ros2bag_data['frames_data']: - frame_data['id'] - for link_data in frame_data['links']: - if link_data['link'] in lost_links: - continue + lost_links = [item for item in topic['links'] if not bpy.data.objects.get(item)] + if lost_links: + logger.warning('Link(s) not found in current scene: %s', lost_links) - parent_location = ((0, 0, 0) if link_data['parent'] == 'world' - else bpy.data.objects[link_data['parent']].location[:]) + for frame_data in topic['frames_data']: + # TODO frame_data['timestamp'] + for link_data in frame_data['links']: + if link_data['link'] in lost_links: + continue - link = bpy.data.objects[link_data['link']] - link.rotation_mode = 'QUATERNION' + parent_location = ((0, 0, 0) if link_data['parent'] == 'world' + else bpy.data.objects[link_data['parent']].location[:]) - link.location = [a + b for a, b in zip(link_data['loc_xyz'], parent_location)] - link.rotation_quaternion = link_data['rot_wxyz'] + link = bpy.data.objects[link_data['link']] + link.rotation_mode = 'QUATERNION' - link.keyframe_insert(data_path='location', frame=frame_data['id']) - link.keyframe_insert(data_path='rotation_quaternion', frame=frame_data['id']) + link.location = [a + b for a, b in zip(link_data['loc_xyz'], parent_location)] + link.rotation_quaternion = link_data['rot_wxyz'] - link.animation_data.action.name = '{}_action'.format(link_data['link']) + link.keyframe_insert(data_path='location', frame=frame_data['id']) + link.keyframe_insert(data_path='rotation_quaternion', frame=frame_data['id']) - logger.info('Setup of %s links is finished!', len(ros2bag_data['scene_links'])) + link.animation_data.action.name = '{}_action'.format(link_data['link']) + + logger.info('Setup of %s links is finished!', len(topic['links'])) + logger.info('Setup of %s topics is finished!', len(ros2bag_data['topics'])) return True 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 deleted file mode 100644 index cfe239c..0000000 --- a/rcg_pipeline/rcg_pipeline/scripts/addons/Robossembler/io_anim_ros2bag/_ros2bag_parser.py +++ /dev/null @@ -1,164 +0,0 @@ -# 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 3fa8ccc..7b250ef 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 @@ -27,81 +27,142 @@ logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) -def append_frame(idx: int, frames, links, frame_data) -> int: +def append_frame(idx: int, frames_data, links, frame_data) -> int: ''' Append datalist for frame and increase frame index. ''' frame_data['links'] = links - frames.append(frame_data) + frames_data.append(frame_data) return idx + 1 -def get_ros2bag_data(ros2bag_path: str, - frame_start=0, frame_end=0, fps=30, target_topic='/tf') -> dict: +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) -> 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] + ros2bag_dir = os.path.dirname(ros2bag_path) - 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 + with AnyReader([Path(ros2bag_dir)]) as ros2bag: + topics = [] + for ros2bag_topic in list(ros2bag.connections): + topic_links = set() + frames_data = [] + # TODO Switch to timestamp instead of frame index + idx = 0 + idx_pre = - 1 + topic_name = None + for connection, timestamp, rawdata in ros2bag.messages(connections=[ros2bag_topic]): + topic_name = connection.topic + 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) + 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) + topic_links.add(t.child_frame_id) + if t.header.frame_id != 'world': + topic_links.add(t.header.frame_id) + link_data = {} + link_data['link'] = t.child_frame_id + link_data['parent'] = t.header.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 + # TODO + if key_link: + idx = append_frame(idx, frames_data, links, frame_data) - assert len(frames_data) - 1 == frames_data[-1]['id'], ( - 'Ros2Bag database has dublicated frames!') + topics.append({ + 'name': topic_name, + 'frame_range': idx, + 'links': list(topic_links), + 'frames_data': frames_data + }) - 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} + return { + 'ros2bag_path': ros2bag_path, + 'frame_start': 0, + 'frame_end': max([topic['topic_frames'] for topic in topics]) - 1, + 'fps': 30, + 'topics': topics + } -#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) +if __name__ == '__main__': + import json + file_path = '/media/disk/robossembler/project/pipeline/projects/subset_0.db3' + json_path = os.path.splitext(file_path)[0] + '.json' + with open(json_path, 'w', encoding='utf-8') as data_file: + json.dump(get_ros2bag_data(file_path), data_file, indent=4) + logger.info('Database saved successfully to %s!', json_path) diff --git a/rcg_pipeline/rcg_pipeline/utils/collection_tools.py b/rcg_pipeline/rcg_pipeline/utils/collection_tools.py index 702579e..d987142 100644 --- a/rcg_pipeline/rcg_pipeline/utils/collection_tools.py +++ b/rcg_pipeline/rcg_pipeline/utils/collection_tools.py @@ -90,3 +90,15 @@ def remove_collections_with_objects(collection=None): bpy.data.objects.remove(obj, do_unlink=True) bpy.data.collections.remove(col) return True + + +def recursive_layer_collection(layer_coll, coll_name): + ''' Set as active collection ''' + 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