def _write_static_data(self, static_group: h5py.Group) -> None:
super()._write_static_data(static_group)
## color and scales of primitive objects
# static_group.create_dataset("target_type", data=self.target_type)
static_group.create_dataset("drop_type", data=self.drop_type)
static_group.create_dataset("drop_position", data=xyz_to_arr(self.drop_position))
static_group.create_dataset("drop_rotation", data=xyz_to_arr(self.drop_rotation))
def _write_static_data(self, static_group: h5py.Group) -> None:
super()._write_static_data(static_group)
# positions and rotations of objects
static_group.create_dataset(
"initial_position",
data=np.stack([xyz_to_arr(p) for p in self.initial_positions], 0))
static_group.create_dataset(
"initial_rotation",
data=np.stack([xyz_to_arr(r) for r in self.initial_rotations], 0))
def _write_static_data(self, static_group: h5py.Group) -> None:
super()._write_static_data(static_group)
## physical
static_group.create_dataset("mass", data=self.masses)
static_group.create_dataset("static_friction",
data=self.static_frictions)
static_group.create_dataset("dynamic_friction",
data=self.dynamic_frictions)
static_group.create_dataset("bounciness", data=self.bouncinesses)
## size and colors
static_group.create_dataset("color", data=self.colors)
static_group.create_dataset(
"scale", data=np.stack([xyz_to_arr(_s) for _s in self.scales], 0))
static_group.create_dataset("scale_x",
data=[_s["x"] for _s in self.scales])
static_group.create_dataset("scale_y",
data=[_s["y"] for _s in self.scales])
static_group.create_dataset("scale_z",
data=[_s["z"] for _s in self.scales])
if self.save_meshes:
mesh_group = static_group.create_group("mesh")
obj_points = []
for idx, object_id in enumerate(self.object_ids):
vertices, faces = self.object_meshes[object_id]
mesh_group.create_dataset(f"faces_{idx}", data=faces)
mesh_group.create_dataset(f"vertices_{idx}", data=vertices)
def _write_static_data(self, static_group: h5py.Group) -> None:
super()._write_static_data(static_group)
## physical
static_group.create_dataset("mass", data=self.masses)
static_group.create_dataset("static_friction", data=self.static_frictions)
static_group.create_dataset("dynamic_friction", data=self.dynamic_frictions)
static_group.create_dataset("bounciness", data=self.bouncinesses)
## size and colors
static_group.create_dataset("color", data=self.colors)
static_group.create_dataset("scale", data=np.stack([xyz_to_arr(_s) for _s in self.scales], 0))
static_group.create_dataset("scale_x", data=[_s["x"] for _s in self.scales])
static_group.create_dataset("scale_y", data=[_s["y"] for _s in self.scales])
static_group.create_dataset("scale_z", data=[_s["z"] for _s in self.scales])
def _place_attachment(self) -> List[dict]:
commands = []
record, data = self.random_primitive(self._attachment_types,
scale=self.attachment_scale_range,
color=self.attachment_color,
exclude_color=self.target_color,
add_data=self.use_attachment)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.attachment = record
self.attachment_id = data['id']
self.attachment_type = data['name']
mass = random.uniform(*get_range(self.attachment_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
if self.attachment_type == 'cylinder':
mass *= (np.pi / 4.0)
elif self.attachment_type == 'cone':
mass *= (np.pi / 12.0)
print("attachment mass", mass)
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": self.tower_height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_attachment))
commands.extend(
self.get_object_material_commands(
record, o_id,
self.get_material_name(self.attachment_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
# for an attachment that is wider at base, better to place links a little higher
a_len, a_height, a_dep = self.get_record_dimensions(record)
if self.attachment_type == 'cone' and self.use_attachment:
self.tower_height += 0.25 * a_height * (
np.sqrt(scale['x']**2 + scale['z']**2) / scale['y'])
if not self.use_attachment:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
# fix it to ground or block
if self.attachment_fixed_to_base and self.use_attachment:
# make it kinematic
if self.use_base:
commands.append({
"$type": "add_fixed_joint",
"parent_id": self.base_id,
"id": o_id
})
elif not self.use_base: # make kinematic
commands.extend([{
"$type": "set_object_collision_detection_mode",
"mode": "continuous_speculative",
"id": o_id
}, {
"$type": "set_kinematic_state",
"id": o_id,
"is_kinematic": True,
"use_gravity": True
}])
# add a cap
if self.use_cap and self.use_attachment:
commands.extend(
self._build_base(height=(self.tower_height +
a_height * scale['y']),
as_cap=True))
if self.attachment_fixed_to_base:
commands.append({
"$type": "add_fixed_joint",
"parent_id": o_id,
"id": self.cap_id
})
return commands
def _build_base(self, height, as_cap=False) -> List[dict]:
commands = []
record, data = self.random_primitive(self._base_types,
scale=self.base_scale_range,
color=self.base_color,
exclude_color=self.target_color,
add_data=self.use_base)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
if as_cap:
self.cap_id = data['id']
self.cap_type = data['name']
else:
self.base_id = data['id']
self.base_type = data['name']
mass = random.uniform(*get_range(self.base_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
print("base mass", mass)
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_base))
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.base_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
if self.use_base:
if self.base_type not in ['pipe', 'torus']:
b_len, b_height, b_dep = self.get_record_dimensions(record)
self.tower_height += b_height * scale['y']
else:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
self.tower_height = 0.0
return commands
def _add_cap(self) -> List[dict]:
commands = []
# the cap becomes the target
record, data = self.random_primitive(self._cap_types,
scale=self.target_scale_range,
color=self.target_color,
add_data=self.use_cap)
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
self.cap = record
self.cap_type = data["name"]
if self.use_cap:
self._replace_target_with_object(record, data)
mass = random.uniform(*get_range(self.middle_mass_range))
mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
mass *= self.STANDARD_MASS_FACTOR
commands.extend(
self.add_physics_object(record=record,
position={
"x": 0.,
"y": self.tower_height,
"z": 0.
},
rotation={
"x": 0.,
"y": 0.,
"z": 0.
},
mass=mass,
dynamic_friction=0.5,
static_friction=0.5,
bounciness=0,
o_id=o_id,
add_data=self.use_cap))
# Set the cap object material
commands.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(self.target_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
if not self.use_cap:
commands.append({
"$type":
self._get_destroy_object_command_name(o_id),
"id":
int(o_id)
})
self.object_ids = self.object_ids[:-1]
else:
self.tower_height += scale["y"]
return commands
def _build_stack(self) -> List[dict]:
commands = []
height = self.tower_height
height += self.zone_scale['y'] if not self.remove_zone else 0.0
# build the block scales
mid = self.num_blocks / 2.0
grad = self.middle_scale_gradient
self.block_scales = [
self._get_block_scale(offset=grad * (mid - i))
for i in range(self.num_blocks)
]
self.blocks = []
# place the blocks
for m in range(self.num_blocks):
record, data = self.random_primitive(
self._middle_types,
scale=self.block_scales[m],
color=self.middle_color,
exclude_color=self.target_color)
self.middle_type = data["name"]
o_id, scale, rgb = [data[k] for k in ["id", "scale", "color"]]
block_pos = self._get_block_position(scale, height)
block_rot = self.get_y_rotation(self.middle_rotation_range)
# scale the mass to give each block a uniform density
block_mass = random.uniform(*get_range(self.middle_mass_range))
block_mass *= (np.prod(xyz_to_arr(scale)) /
np.prod(xyz_to_arr(self.STANDARD_BLOCK_SCALE)))
block_mass *= self.STANDARD_MASS_FACTOR
commands.extend(
self.add_physics_object(record=record,
position=block_pos,
rotation=block_rot,
mass=block_mass,
dynamic_friction=random.uniform(
0, 0.9),
static_friction=random.uniform(0, 0.9),
bounciness=random.uniform(0, 1),
o_id=o_id))
# Set the block object material
commands.extend(
self.get_object_material_commands(
record, o_id,
self.get_material_name(self.middle_material)))
# Scale the object and set its color.
commands.extend([{
"$type": "set_color",
"color": {
"r": rgb[0],
"g": rgb[1],
"b": rgb[2],
"a": 1.
},
"id": o_id
}, {
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
}])
print("placed middle object %s" % str(m + 1))
# update height
_y = record.bounds['top']['y'] if self.middle_type != 'bowl' else (
record.bounds['bottom']['y'] + 0.1)
height += scale["y"] * _y
data.update({
'position': block_pos,
'rotation': block_rot,
'mass': block_mass
})
print("middle object data", data)
self.blocks.append((record, data))
self.tower_height = height
return commands
def add_primitive(self,
record: ModelRecord,
position: Dict[str, float] = TDWUtils.VECTOR3_ZERO,
rotation: Dict[str, float] = TDWUtils.VECTOR3_ZERO,
scale: Dict[str, float] = {
"x": 1.,
"y": 1.,
"z": 1
},
o_id: Optional[int] = None,
material: Optional[str] = None,
color: Optional[list] = None,
exclude_color: Optional[list] = None,
mass: Optional[float] = 2.0,
dynamic_friction: Optional[float] = 0.1,
static_friction: Optional[float] = 0.1,
bounciness: Optional[float] = 0,
add_data: Optional[bool] = True,
scale_mass: Optional[bool] = True,
make_kinematic: Optional[bool] = False,
obj_list: Optional[list] = [],
apply_texture: Optional[bool] = True) -> List[dict]:
cmds = []
if o_id is None:
o_id = self._get_next_object_id()
if scale_mass:
mass = mass * np.prod(xyz_to_arr(scale))
# add the physics stuff
cmds.extend(
self.add_physics_object(record=record,
position=position,
rotation=rotation,
mass=mass,
dynamic_friction=dynamic_friction,
static_friction=static_friction,
bounciness=bounciness,
o_id=o_id,
add_data=add_data))
# scale the object
cmds.append({
"$type": "scale_object",
"scale_factor": scale,
"id": o_id
})
# set the material and color
if apply_texture:
cmds.extend(
self.get_object_material_commands(
record, o_id, self.get_material_name(material)))
color = color if color is not None else self.random_color(
exclude=exclude_color)
cmds.append({
"$type": "set_color",
"color": {
"r": color[0],
"g": color[1],
"b": color[2],
"a": 1.
},
"id": o_id
})
if make_kinematic:
cmds.extend([{
"$type": "set_object_collision_detection_mode",
"mode": "continuous_speculative",
"id": o_id
}, {
"$type": "set_kinematic_state",
"id": o_id,
"is_kinematic": True,
"use_gravity": True
}])
if add_data:
data = {
'name': record.name,
'id': o_id,
'scale': scale,
'color': color,
'material': material,
'mass': mass,
'dynamic_friction': dynamic_friction,
'static_friction': static_friction,
'bounciness': bounciness
}
self._add_name_scale_color(record, data)
obj_list.append((record, data))
return cmds