Blender Addon: import ros2bag animation and assign to blender

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),
+    'version': (0, 2),
-    'blender': (3, 6, 1),
+    'blender': (4, 2, 0),
     'location': '3D View > Toolbox',
     'description': 'Robossembler pipeline tools',
     'warning': '',
@@ -77,10 +78,10 @@ class addon_Properties(PropertyGroup):
     )
 
 
-class RobossemblerPanel(Panel):
+class RobossemblerPanel1(Panel):
-    '''Doc'''
+    ''' Robossembler UI'''
-    bl_label = 'Robossembler Tools'
+    bl_idname = 'ROBOSSEMBLER_PT_EXPORT_JSON'
-    bl_idname = 'ROBOSSEMBLER_PT_PANEL'
+    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,
+    RobossemblerPanel1,
-    RobossemblerOperator,
+    RobossemblerPanel2,
+    RobossemblerOperator1,
+    RobossemblerOperator2,
     addon_Properties)
 
 

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 <http://www.gnu.org/licenses/>.
+'''
+
+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

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 <http://www.gnu.org/licenses/>.
+'''
+
+import collections
+import logging
+import os
+from pathlib import Path
+from rosbags.highlevel import AnyReader
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def append_frame(idx: int, frames, links, frame_data) -> int:
+    ''' Append datalist for frame and increase frame index. '''
+    frame_data['links'] = links
+    frames.append(frame_data)
+    return idx + 1
+
+
+def gimbal_locks_filter(ros2bag_func):
+    ''' Finding and fix gimbal locks. '''
+    def wrapper(*args, **kwargs):
+        # finding gimbal locks
+        ros2bag_data = ros2bag_func(*args, **kwargs)
+        gimbal_locks = {}
+        for link in ros2bag_data['scene_links']:
+            gimbal_data = collections.defaultdict(list)
+            link_rotation = []
+            for frame_data in ros2bag_data['frames_data']:
+                link_rotation_pre = []
+                for link_data in frame_data['links']:
+                    if link_data['link'] == link:
+                        link_rotation_pre = link_rotation
+                        link_rotation = link_data['rot_wxyz']
+                if link_rotation_pre:
+                    idx = 0
+                    for axis_pre, axis_cur in zip(link_rotation_pre, link_rotation):
+                        if (round(axis_cur, 2) > 0 > round(axis_pre, 2)
+                                or round(axis_cur, 2) < 0 < round(axis_pre, 2)):
+                            gimbal_data[frame_data['id']].append(idx)
+                            gimbal_locks[link] = gimbal_data
+                            logger.info(
+                                'Gimbal lock detected for %s link %s axis at %s frame',
+                                link, idx, frame_data['id'])
+                        idx += 1
+
+        # fix gimbal locks
+        for link in gimbal_locks:
+            gimbal_frame = list(gimbal_locks[link])[0]
+            if len(gimbal_locks[link]) == 2:
+                gimbal_range = list(gimbal_locks[link])
+            else:
+                at_start = abs(ros2bag_data['frame_start'] - gimbal_frame)
+                at_end = abs(ros2bag_data['frame_end'] - gimbal_frame)
+                if at_start < at_end:
+                    gimbal_range = [ros2bag_data['frame_start'], gimbal_frame]
+                else:
+                    gimbal_range = [gimbal_frame, ros2bag_data['frame_end']]
+            for frame_data in ros2bag_data['frames_data'][gimbal_range[0]:gimbal_range[1]]:
+                for link_data in frame_data['links']:
+                    if link_data['link'] == link:
+                        for axis in gimbal_locks[link][gimbal_frame]:
+                            link_data['rot_wxyz'][axis] = (
+                                -1 * link_data['rot_wxyz'][axis])
+                            logger.debug(
+                                'Gimbal lock fixed for %s link %s axis %s',
+                                link, axis, frame_data['id'])
+            logger.info('Gimbal lock fixed for %s link', link)
+
+        return ros2bag_data
+
+    return wrapper
+
+
+@gimbal_locks_filter
+def get_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)

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):