robot builder wip
This commit is contained in:
brothermechanic
parent
ba6a8e9d08
commit
bafa332cac
10 changed files with 381 additions and 256 deletions
|
|
@ -31,5 +31,5 @@ import rcg_pipeline
|
|||
project_dir = '/path/to/<my_project_dir>'
|
||||
|
||||
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)
|
||||
```
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
127
rcg_pipeline/rcg_pipeline/robot_asset.py
Normal file
127
rcg_pipeline/rcg_pipeline/robot_asset.py
Normal file
|
|
@ -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 <https://www.gnu.org/licenses/>.
|
||||
#
|
||||
# ***** 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
|
||||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -28,22 +28,24 @@ 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)]
|
||||
for topic in ros2bag_data['topics']:
|
||||
# TODO if topic['name']
|
||||
|
||||
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)
|
||||
|
||||
for frame_data in ros2bag_data['frames_data']:
|
||||
frame_data['id']
|
||||
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
|
||||
|
|
@ -62,5 +64,6 @@ def set_animation_data(context, ros2bag_path, frame_start=0, frame_end=0, fps=30
|
|||
|
||||
link.animation_data.action.name = '{}_action'.format(link_data['link'])
|
||||
|
||||
logger.info('Setup of %s links is finished!', len(ros2bag_data['scene_links']))
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 <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_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)
|
||||
|
|
@ -27,31 +27,88 @@ 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]
|
||||
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
|
||||
for connection, timestamp, rawdata in ros2bag.messages(connections=targets):
|
||||
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 = []
|
||||
|
|
@ -65,12 +122,12 @@ def get_ros2bag_data(ros2bag_path: str,
|
|||
idx = append_frame(idx, frames_data, links, frame_data)
|
||||
else:
|
||||
key_link.append(c_key)
|
||||
scene_links.add(t.child_frame_id)
|
||||
topic_links.add(t.child_frame_id)
|
||||
if t.header.frame_id != 'world':
|
||||
scene_links.add(t.header.frame_id)
|
||||
topic_links.add(t.header.frame_id)
|
||||
link_data = {}
|
||||
link_data['parent'] = t.header.frame_id
|
||||
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,
|
||||
|
|
@ -82,26 +139,30 @@ def get_ros2bag_data(ros2bag_path: str,
|
|||
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__':
|
||||
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)
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue