From bafa332cacc8a4fd7213a1ca8f67e1f82e0830b7 Mon Sep 17 00:00:00 2001 From: brothermechanic Date: Mon, 1 Jul 2024 14:12:33 +0300 Subject: [PATCH] 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