From db3e0d127358e52530adfe7b2d0fc53d5073b1b0 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Sat, 24 Feb 2024 14:11:29 +0300 Subject: [PATCH] Blender Addon: import ros2bag animation and assign to blender --- .../scripts/addons/Robossembler/__init__.py | 69 ++++++-- .../Robossembler/io_anim_ros2bag/__init__.py | 66 +++++++ .../io_anim_ros2bag/ros2bag_parser.py | 163 ++++++++++++++++++ .../__init__.py} | 2 +- .../{ => io_scene_json}/model_md5.py | 0 .../{ => io_scene_json}/model_name.py | 0 .../{ => io_scene_json}/model_paths.py | 0 7 files changed, 285 insertions(+), 15 deletions(-) create mode 100644 cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py create mode 100644 cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py rename cg/blender/scripts/addons/Robossembler/{scene_to_json.py => io_scene_json/__init__.py} (99%) rename cg/blender/scripts/addons/Robossembler/{ => io_scene_json}/model_md5.py (100%) rename cg/blender/scripts/addons/Robossembler/{ => io_scene_json}/model_name.py (100%) rename cg/blender/scripts/addons/Robossembler/{ => io_scene_json}/model_paths.py (100%) diff --git a/cg/blender/scripts/addons/Robossembler/__init__.py b/cg/blender/scripts/addons/Robossembler/__init__.py index 2b52d89..0f46f8f 100644 --- a/cg/blender/scripts/addons/Robossembler/__init__.py +++ b/cg/blender/scripts/addons/Robossembler/__init__.py @@ -40,13 +40,14 @@ from bpy.props import ( EnumProperty, PointerProperty) -from .scene_to_json import export_scene_conf +from .io_scene_json import export_json +from .io_anim_ros2bag import set_animation_data bl_info = { 'name': 'Robossembler Tools', 'author': 'brothermechanic@gmail.com', - 'version': (0, 1), - 'blender': (3, 6, 1), + 'version': (0, 2), + 'blender': (4, 2, 0), 'location': '3D View > Toolbox', 'description': 'Robossembler pipeline tools', 'warning': '', @@ -77,10 +78,10 @@ class addon_Properties(PropertyGroup): ) -class RobossemblerPanel(Panel): - '''Doc''' - bl_label = 'Robossembler Tools' - bl_idname = 'ROBOSSEMBLER_PT_PANEL' +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' @@ -97,29 +98,69 @@ class RobossemblerPanel(Panel): col.scale_y = 2.0 col.operator('export.scene_config', icon='WORLD_DATA', - text='Export Scene Config') + text='Export Scene') -class RobossemblerOperator(Operator): +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, 'conf_file_path') + + col = layout.column() + col.alert = True + col.scale_y = 2.0 + col.operator('export.scene_config', + icon='ACTION', + text='Import Animation') + + +class RobossemblerOperator1(Operator): '''Tooltip''' bl_idname = 'export.scene_config' bl_label = '' - bl_description = 'Export scene config' + bl_description = 'Export scene liks to json config.' bl_options = {'REGISTER', 'UNDO'} def execute(self, context): prop = context.scene.robossembler_properties - conf_file_path = os.path.realpath(bpy.path.abspath((prop.conf_file_path))) + file_path = os.path.realpath(bpy.path.abspath((prop.conf_file_path))) physics_engine = prop.engine - export_scene_conf(context, conf_file_path, physics_engine) + export_json(context, file_path, physics_engine) + + return {'FINISHED'} + + +class RobossemblerOperator2(Operator): + '''Tooltip''' + bl_idname = 'export.scene_config' + bl_label = '' + bl_description = 'Export scene liks to json config.' + bl_options = {'REGISTER', 'UNDO'} + + def execute(self, context): + prop = context.scene.robossembler_properties + + file_path = os.path.realpath(bpy.path.abspath((prop.conf_file_path))) + set_animation_data(context, file_path) return {'FINISHED'} classes = ( - RobossemblerPanel, - RobossemblerOperator, + RobossemblerPanel1, + RobossemblerPanel2, + RobossemblerOperator1, + RobossemblerOperator2, addon_Properties) diff --git a/cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py b/cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/__init__.py new file mode 100644 index 0000000..e94b048 --- /dev/null +++ b/cg/blender/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/cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py b/cg/blender/scripts/addons/Robossembler/io_anim_ros2bag/ros2bag_parser.py new file mode 100644 index 0000000..7908d3b --- /dev/null +++ b/cg/blender/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/cg/blender/scripts/addons/Robossembler/scene_to_json.py b/cg/blender/scripts/addons/Robossembler/io_scene_json/__init__.py similarity index 99% rename from cg/blender/scripts/addons/Robossembler/scene_to_json.py rename to cg/blender/scripts/addons/Robossembler/io_scene_json/__init__.py index 8bc242d..732f24b 100644 --- a/cg/blender/scripts/addons/Robossembler/scene_to_json.py +++ b/cg/blender/scripts/addons/Robossembler/io_scene_json/__init__.py @@ -47,7 +47,7 @@ logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) -def export_scene_conf(context, scene_config_path, physics_engine): +def export_json(context, scene_config_path, physics_engine): ''' Export scene config to json. ''' # open conf file if os.path.isfile(scene_config_path): diff --git a/cg/blender/scripts/addons/Robossembler/model_md5.py b/cg/blender/scripts/addons/Robossembler/io_scene_json/model_md5.py similarity index 100% rename from cg/blender/scripts/addons/Robossembler/model_md5.py rename to cg/blender/scripts/addons/Robossembler/io_scene_json/model_md5.py diff --git a/cg/blender/scripts/addons/Robossembler/model_name.py b/cg/blender/scripts/addons/Robossembler/io_scene_json/model_name.py similarity index 100% rename from cg/blender/scripts/addons/Robossembler/model_name.py rename to cg/blender/scripts/addons/Robossembler/io_scene_json/model_name.py diff --git a/cg/blender/scripts/addons/Robossembler/model_paths.py b/cg/blender/scripts/addons/Robossembler/io_scene_json/model_paths.py similarity index 100% rename from cg/blender/scripts/addons/Robossembler/model_paths.py rename to cg/blender/scripts/addons/Robossembler/io_scene_json/model_paths.py