175 lines
5.8 KiB
Python
175 lines
5.8 KiB
Python
import dataclasses
|
|
|
|
from flask import jsonify, request
|
|
|
|
import gazebo_ctrl
|
|
import scene_plugin
|
|
|
|
from .gen_models import GraspModelGenerator
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class BoxModel:
|
|
length: float
|
|
weight: float
|
|
height: float
|
|
tag_id: int
|
|
mass: float
|
|
type: str = 'box'
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TableModel:
|
|
height: float
|
|
type: str = 'table'
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TagModel:
|
|
tag_id: int
|
|
type: str = 'tag'
|
|
|
|
|
|
class Plugin(scene_plugin.ScenePluginBase):
|
|
def __init__(self, scene_name):
|
|
super().__init__(scene_name)
|
|
self.model_dict = {}
|
|
self.spawner = GraspModelGenerator()
|
|
self.gz_ctrl = gazebo_ctrl.GazeboROSController()
|
|
|
|
def post_scene_model_spawn(self):
|
|
print(222)
|
|
try:
|
|
# 解析参数
|
|
data = request.get_json(force=True)
|
|
name = data.get('name')
|
|
type_ = data.get('type')
|
|
pose = data.get('pose')
|
|
position = pose.get('position')
|
|
orientation = pose.get('orientation')
|
|
tag_id = data.get('tag_id', 0)
|
|
mass = data.get('mass', 1)
|
|
size = data.get('size')
|
|
if not name or not type_ or pose is None:
|
|
raise ValueError('name, type, and pose are required')
|
|
# 生成模型
|
|
if type_ == 'box':
|
|
if not size or not all(k in size for k in ('l', 'h', 'w')):
|
|
raise ValueError('size with l, h, w is required for box')
|
|
model = BoxModel(
|
|
length=size['l'],
|
|
height=size['h'],
|
|
weight=size['w'],
|
|
tag_id=tag_id,
|
|
mass=mass
|
|
)
|
|
ret, msg = self.spawner.generate_box_with_apriltag(
|
|
name,
|
|
l=size['l'],
|
|
w=size['w'],
|
|
h=size['h'],
|
|
mass=mass,
|
|
tag_id=tag_id,
|
|
x=position.get('x', 0),
|
|
y=position.get('y', 0),
|
|
z=position.get('z', 0),
|
|
roll=orientation.get('roll', 0),
|
|
pitch=orientation.get('pitch', 0),
|
|
yaw=orientation.get('yaw', 0)
|
|
)
|
|
if ret == False:
|
|
raise RuntimeError(msg)
|
|
elif type_ == 'table':
|
|
model = TableModel(
|
|
height=pose.get('z', 1)
|
|
)
|
|
self.spawner.generate_table(
|
|
name,
|
|
x=position.get('x', 0),
|
|
y=position.get('y', 0),
|
|
height=position.get('z', 0),
|
|
roll=orientation.get('roll', 0),
|
|
pitch=orientation.get('pitch', 0),
|
|
yaw=orientation.get('yaw', 0)
|
|
)
|
|
elif type_ == 'tag':
|
|
model = TagModel(
|
|
tag_id=tag_id
|
|
)
|
|
self.spawner.generate_apriltag(
|
|
name,
|
|
tag_id=tag_id,
|
|
x=position.get('x', 0),
|
|
y=position.get('y', 0),
|
|
z=position.get('z', 0),
|
|
roll=orientation.get('roll', 0),
|
|
pitch=orientation.get('pitch', 0),
|
|
yaw=orientation.get('yaw', 0)
|
|
)
|
|
else:
|
|
raise ValueError('Invalid type')
|
|
# 注册模型
|
|
self.model_dict[name] = model
|
|
except Exception as e:
|
|
print(str(e))
|
|
return str(e), 400
|
|
return 'OK', 200
|
|
|
|
def get_scene_model_state(self):
|
|
try:
|
|
res = self.gz_ctrl.get_model_state_f(self.model_dict.keys())
|
|
for item in res:
|
|
model = self.model_dict.get(item['name'])
|
|
if model.type == 'box':
|
|
item['tag_id'] = model.tag_id
|
|
item['mass'] = model.mass
|
|
item['size'] = {'l': model.length, 'w': model.weight, 'h': model.height}
|
|
elif model.type == 'table':
|
|
item['pose']['position']['z'] = model.height
|
|
elif model.type == 'tag':
|
|
item['tag_id'] = model.tag_id
|
|
return jsonify(res), 200
|
|
except Exception as e:
|
|
return str(e), 400
|
|
|
|
def post_scene_model_state(self):
|
|
try:
|
|
# 解析参数
|
|
data = request.get_json(force=True)
|
|
name = data.get('name')
|
|
type_ = data.get('type')
|
|
pose = data.get('pose')
|
|
position = pose.get('position')
|
|
orientation = pose.get('orientation')
|
|
tag_id = data.get('tag_id', 0)
|
|
mass = data.get('mass', 0)
|
|
size = data.get('size')
|
|
if not name or not type_ or pose is None:
|
|
raise ValueError('name, type, and pose are required')
|
|
# 修改模型
|
|
model = self.model_dict[name]
|
|
if model.type != type_:
|
|
raise ValueError(f'Cannot change model type from {self.model_dict[name].type} to {type_}')
|
|
# 修改pose
|
|
pass
|
|
# 修改独特属性
|
|
if type_ == 'box':
|
|
if mass:
|
|
model.mass = mass
|
|
if tag_id:
|
|
model.tag_id = tag_id
|
|
elif type_ == 'table':
|
|
height = position.get('z', 0)
|
|
if height:
|
|
model.height = height
|
|
elif type_ == 'tag':
|
|
if tag_id:
|
|
model.tag_id = tag_id
|
|
else:
|
|
raise ValueError('Invalid type')
|
|
except Exception as e:
|
|
return str(e), 400
|
|
return 'OK', 200
|
|
|
|
if __name__ == '__main__':
|
|
pass |