def __init__(self):
self.obj_name_dict: Dict[int, str] = {}
# Set up the object transform data.
object_setup_data = json.loads(Path("object_setup.json").read_text())
self.object_setups: List[_ObjectSetup] = []
for o in object_setup_data:
combo = _ObjectSetup(id=int(o),
model_name=object_setup_data[o]["model_name"],
position=object_setup_data[o]["position"],
rotation=object_setup_data[o]["rotation"],
scale=object_setup_data[o]["scale"])
self.object_setups.append(combo)
# Parse the default objects.csv spreadsheet.
self.object_audio_data = PyImpact.get_object_info()
self.temp_amp = 0.1
# Keep track of the current trial number, for logging purposes.
self.current_trial_num = 0
# Fetch the ball and board model's records; we will need them later to change its material.
self.special_models = ModelLibrarian(library="models_special.json")
self.full_models = ModelLibrarian(library="models_full.json")
self.ball_record = self.special_models.get_record("prim_sphere")
self.board_record = self.full_models.get_record("wood_board")
# Set path to write out logging info.
self.root_dest_dir = Path("dist/mode_properties_logs")
if not self.root_dest_dir.exists():
self.root_dest_dir.mkdir(parents=True)
super().__init__()
def __init__(self):
self.model_list = [
"b03_db_apps_tech_08_04", "trashbin", "trunck",
"whirlpool_akzm7630ix", "satiro_sculpture", "towel-radiator-2",
"b03_folding-screen-panel-room-divider",
"naughtone_pinch_stool_chair", "microwave", "trunk_6810-0009",
"suitcase", "kayak_small", "elephant_bowl", "trapezoidal_table",
"b05_pc-computer-printer-1", "dishwasher_4",
"chista_slice_of_teak_table", "buddah", "b05_elsafe_infinity_ii",
"backpack", "b06_firehydrant_lod0", "b05_ticketmachine",
"b05_trophy", "b05_kitchen_aid_toster", "b05_heavybag",
"bongo_drum_hr_blend", "b03_worldglobe", "ceramic_pot",
"b04_kenmore_refr_70419", "b03_zebra", "b05_gibson_j-45",
"b03_cow", "b03_sheep", "b04_stringer"
]
# Cache the record for the receptacle.
self.receptacle_record = ModelLibrarian(
"models_special.json").get_record("fluid_receptacle1x1")
# Cache the fluid types.
self.ft = FluidTypes()
self.full_lib = ModelLibrarian("models_full.json")
self.pool_id = None
super().__init__()
def __init__(self, port: int = 1071):
super().__init__(port=port)
self._occluders: List[ModelRecord] = ModelLibrarian(str(Path("occluders.json").resolve())).records
self._ball = ModelLibrarian("models_flex.json").get_record("sphere")
self._ball_id = 0
self._occ_id = 1
self.material_librarian = MaterialLibrarian()
def __init__(self, port: int = 1071):
super().__init__(port=port)
self.ball_record = ModelLibrarian("models_flex.json").get_record(
"sphere")
if sys.argv[2] == '30':
self.ramp_record = ModelLibrarian("models_full.json").get_record(
"ramp_with_platform_30")
elif sys.argv[2] == '60':
self.ramp_record = ModelLibrarian("models_full.json").get_record(
"ramp_with_platform_60")
self.pin_record = ModelLibrarian("models_flex.json").get_record(
"cylinder")
self.ball_material = "marble_white"
self.start()
def _get_librarian(self, description: str) -> ModelLibrarian:
"""
Returns a librarian object.
:param description: The description of the library.
"""
ModelLibrarian.create_library(description, self.library_path)
print("Adding records to the library...")
return ModelLibrarian(str(self.library_path.resolve()))
def __init__(self, port: int = 1071):
# Load the objects.
self.object_records = ModelLibrarian(str(
Path("flex.json").resolve())).records
# Get the cloth record.
self.cloth_record = MODEL_LIBRARIES["models_special.json"].get_record(
"cloth_square")
self.cloth_id = 0
super().__init__(port=port)
def create_asset_bundles(self,
batch_size: int = 1000,
vhacd_resolution: int = 8000000,
first_batch_only: bool = False) -> None:
"""
Convert all .obj files into asset bundles.
:param batch_size: The number of models per batch.
:param vhacd_resolution: Higher value=better-fitting colliders and slower build process.
:param first_batch_only: If true, output only the first batch. Useful for testing purposes.
"""
records = ModelLibrarian(
library=str(self.library_path.resolve())).records
a = AssetBundleCreator(quiet=True)
pbar = tqdm(total=len(records))
while len(records) > 0:
# Get the next batch.
batch: List[ModelRecord] = records[:batch_size]
records = records[batch_size:]
for record in batch:
# If the asset bundles for this record already exist, skip it.
urls_exist = False
for platform in record.urls:
url = record.urls[platform][8:]
if Path(url).exists():
urls_exist = True
if urls_exist:
continue
# If the prefab for this record exists, skip it.
dest_path = Path.home().joinpath(
f"asset_bundle_creator/Assets/Resources/models/{record.name}.obj"
)
if dest_path.exists():
continue
# Process the .obj
obj_path = self._get_obj(record)
# Move the files and remove junk.
a.move_files_to_unity_project(
None,
model_path=obj_path,
sub_directory=f"models/{record.name}")
# Creating the asset bundles.
a.create_many_asset_bundles(str(self.library_path.resolve()),
cleanup=True,
vhacd_resolution=vhacd_resolution)
pbar.update(len(batch))
# Process only the first batch of models.
if first_batch_only:
break
pbar.close()
def __init__(self, port: int = 1071):
self.toy_records = ModelLibrarian(str(
Path("toys.json").resolve())).records
self.ramp_positions = [{
"x": 3.5,
"y": 0.02,
"z": 1.5
}, {
"x": -1,
"y": 0.02,
"z": 2.38
}, {
"x": 4.58,
"y": 0.02,
"z": -2.85
}, {
"x": -0.94,
"y": 0.02,
"z": -2.83
}, {
"x": -3.4,
"y": 0.02,
"z": 0
}]
self.ramp_rotations = [{
"x": 0,
"y": -45,
"z": 0
}, {
"x": 0,
"y": -90,
"z": 0
}, {
"x": 0,
"y": 30,
"z": 0
}, {
"x": -90,
"y": -190,
"z": 0
}, {
"x": -90,
"y": -120,
"z": 0
}, {
"x": 0,
"y": 120,
"z": 0
}]
super().__init__(port=port)
def run(self):
self.start()
self.communicate(TDWUtils.create_empty_room(20, 20))
self.communicate(TDWUtils.create_avatar(position={"x": 0, "y": 3, "z": -6},
look_at=TDWUtils.VECTOR3_ZERO))
model_name = "rh10"
z = -3
x = -1.5
print("With the add_object command (complex syntax but you have maximum control):")
record = ModelLibrarian().get_record(model_name)
self.communicate({"$type": "add_object",
"name": model_name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": {"x": x, "y": 0, "z": z},
"rotation": TDWUtils.VECTOR3_ZERO,
"category": record.wcategory,
"id": self.get_unique_id()})
x = 0
print("With the wrapper function Controller.add_object() "
"(easy to use, but you can't add additional commands to this frame):")
self.add_object(model_name=model_name,
position={"x": x, "y": 0, "z": z},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_core.json")
x = 1.5
print("With the wrapper function Controller.get_add_object() "
"(harder to use, but you can add commands to this frame):")
self.communicate(self.get_add_object(model_name=model_name,
object_id=self.get_unique_id(),
position={"x": x, "y": 0, "z": 0},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_core.json"))
print("With the add_object command, minus all optional parameters (the model won't scale properly!):")
self.communicate({"$type": "add_object",
"name": model_name,
"url": record.get_url(),
"id": self.get_unique_id()})
print("With the wrapper function Controller.add_object(), minus all optional parameters:")
self.add_object(model_name)
print("With the wrapper function Controller.get_add_object(), minus all optional parameters:")
self.communicate(self.get_add_object(model_name=model_name,
object_id=self.get_unique_id()))
def __init__(self):
self.obj_name_dict: Dict[int, str] = {}
# Set up the object transform data.
object_setup_data = json.loads(Path("object_setup.json").read_text())
self.object_setups: List[_ObjectSetup] = []
for o in object_setup_data:
combo = _ObjectSetup(id=int(o),
model_name=object_setup_data[o]["model_name"],
position=object_setup_data[o]["position"],
rotation=object_setup_data[o]["rotation"],
scale=object_setup_data[o]["scale"])
self.object_setups.append(combo)
# Parse the default objects.csv spreadsheet.
self.object_audio_data = PyImpact.get_object_info()
# Fetch the ball and board model's records; we will need them later to change its material.
self.special_models = ModelLibrarian(library="models_special.json")
self.full_models = ModelLibrarian(library="models_full.json")
self.ball_record = self.special_models.get_record("prim_sphere")
self.board_record = self.full_models.get_record("wood_board")
super().__init__()
def __init__(self, port: int = 1071):
super().__init__(port=port)
# Cache the ball data.
self._ball = ModelLibrarian("models_special.json").get_record(
"prim_sphere")
self._ball_id = 0
# The position the ball starts in and the position the ball is directed at.
self._p0: Dict[str, float] = {}
self._p1: Dict[str, float] = {}
# Cache the skybox records.
skybox_lib = HDRISkyboxLibrarian()
self._skyboxes: List[str] = [
r.name for r in skybox_lib.records if r.sun_intensity >= 0.8
]
def _initialize_scene(self, c: Controller) -> List[dict]:
c.model_librarian = ModelLibrarian("models_full.json")
# Initialize the scene.
commands = super()._initialize_scene(c)
chair_name = "brown_leather_dining_chair"
# Create the the table.
commands.extend(self._init_object(c=c,
name="quatre_dining_table",
pos=TDWUtils.VECTOR3_ZERO,
rot=TDWUtils.VECTOR3_ZERO))
# Create 8 chairs around the table.
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": 0, "y": 0, "z": -1.55},
rot={"x": 0, "y": 0, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": 0, "y": 0, "z": 1.55},
rot={"x": 0, "y": 180, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": -1, "y": 0, "z": -0.85},
rot={"x": 0, "y": 90, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": -1, "y": 0, "z": 0},
rot={"x": 0, "y": 90, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": -1, "y": 0, "z": 0.85},
rot={"x": 0, "y": 90, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": 1, "y": 0, "z": -0.85},
rot={"x": 0, "y": -90, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": 1, "y": 0, "z": 0},
rot={"x": 0, "y": -90, "z": 0}))
commands.extend(self._init_object(c=c,
name=chair_name,
pos={"x": 1, "y": 0, "z": 0.85},
rot={"x": 0, "y": -90, "z": 0}))
return commands
def __init__(self, port: int = 1071):
super().__init__(port=port)
self.model_librarian = ModelLibrarian("models_flex.json")
# A list of functions that will return commands to initialize a trial.
self.scenarios = [
self.drop_onto_floor, self.drop_onto_object, self.throw_into_wall,
self.push_into_other
]
# Size of the room bounds (used for placing objects near walls).
self.env_bounds = {"x": 4, "y": 0, "z": 2}
# Cached pressure parameters per model.
self.pressures = loads(Path("squish_pressures.json").read_text())
# Only use records in the pressures dictionary.
self.records = [
r for r in self.model_librarian.records if r.name in self.pressures
]
def get_add_object(self,
model_name: str,
object_id: int,
position={
"x": 0,
"y": 0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library: str = "") -> dict:
"""
Returns a valid add_object command.
:param model_name: The name of the model.
:param position: The position of the model.
:param rotation: The starting rotation of the model, in Euler angles.
:param library: The path to the records file. If left empty, the default library will be selected. See `ModelLibrarian.get_library_filenames()` and `ModelLibrarian.get_default_library()`.
:param object_id: The ID of the new object.
:return An add_object command that the controller can then send.
"""
if self.model_librarian is None or (
library != "" and self.model_librarian.library != library):
self.model_librarian = ModelLibrarian(library=library)
record = self.model_librarian.get_record(model_name)
return {
"$type": "add_object",
"name": model_name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": position,
"rotation": rotation,
"category": record.wcategory,
"id": object_id
}
def run(c: Controller) -> None:
"""
Check every model for missing materials.
:param c: The controller.
"""
# Create a new output file.
if MissingMaterials.OUTPUT_FILE.exists():
MissingMaterials.OUTPUT_FILE.unlink()
MissingMaterials.OUTPUT_FILE = str(
MissingMaterials.OUTPUT_FILE.resolve())
MissingMaterials.start(c)
print(
f"The names of models with missing materials will be saved to: {MissingMaterials.OUTPUT_FILE}"
)
for library_path in ModelLibrarian.get_library_filenames():
print(library_path)
lib = ModelLibrarian(library=library_path)
pbar = tqdm(total=len(lib.records))
for record in lib.records:
if record.do_not_use:
pbar.update(1)
continue
pbar.set_description(record.name)
# Check for missing materials.
if MissingMaterials.materials_are_missing(
c, record, record.get_url()):
with io.open(MissingMaterials.OUTPUT_FILE, "at") as f:
f.write("\n" + record.name)
# Cleanup.
c.communicate([{
"$type": "destroy_object",
"id": MissingMaterials.OBJECT_ID
}, {
"$type": "unload_asset_bundles"
}])
pbar.update(1)
pbar.close()
c.communicate({"$type": "terminate"})
def run(self):
# Parse the substructure of fridge_large.
record = ModelLibrarian().get_record("fridge_large")
self.start()
self.communicate(TDWUtils.create_empty_room(12, 12))
# Create the object.
o_id = self.add_object("fridge_large")
# Add an avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": 0,
"y": 1.886,
"z": -0.15
},
look_at=TDWUtils.VECTOR3_ZERO))
# Disable post-processing, so the changes to the material aren't blurry.
self.communicate({"$type": "set_post_process", "value": False})
for quality in ["low", "med", "high"]:
# Set the visual materials.
self.communicate(
TDWUtils.set_visual_material(self,
record.substructure,
o_id,
"parquet_long_horizontal_clean",
quality=quality))
sleep(5)
self.communicate({
"$type": "unload_asset_bundles",
"bundle_type": "materials"
})
self.communicate({
"$type": "unload_asset_bundles",
"bundle_type": "models"
})
def _initialize_scene(self, c: Controller) -> List[dict]:
c.model_librarian = ModelLibrarian("models_full.json")
commands = super()._initialize_scene(c)
# Add a table, chair, and boxes.
commands.extend(self._init_object(c, "b05_table_new",
pos={"x": 0, "y": 0, "z": 4.33},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "chair_willisau_riale",
pos={"x": 0, "y": 0, "z": 3.7},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "iron_box",
pos={"x": 0.13, "y": 0.65, "z": 4.83},
rot=TDWUtils.VECTOR3_ZERO))
commands.extend(self._init_object(c, "iron_box",
pos={"x": -0.285, "y": 1.342, "z": 4.79},
rot={"x": 90, "y": 0, "z": 0}))
# Add a shelf with a custom scale.
shelf_id = c.get_unique_id()
shelf_name = "metal_lab_shelf"
Scene.OBJECT_IDS.update({shelf_id: shelf_name})
commands.extend([c.get_add_object(shelf_name,
object_id=shelf_id,
rotation={"x": 0, "y": -90, "z": 0},
position={"x": 0, "y": 0, "z": 4.93}),
{"$type": "set_mass",
"id": shelf_id,
"mass": 400},
{"$type": "set_physic_material",
"id": shelf_id,
"bounciness": Scene._OBJECT_INFO[shelf_name].bounciness,
"static_friction": 0.1,
"dynamic_friction": 0.8},
{"$type": "scale_object",
"id": shelf_id,
"scale_factor": {"x": 1, "y": 1.5, "z": 1.8}}])
return commands
def run(self):
""" Generate room using COCO_TDW dataset """
objects_in_scene = 15
object_ids = []
# Get Category-Object mapping
TDW_COCO_models = TDW_relationships.get_COCO_TDW_mapping()
# print("TDWCOCO:", TDW_COCO_models)
# print("KEYS:", TDW_COCO_models.keys())
# Gets COCO categories co-occurring in a scene
# +5 is for dealing with failed object insertion attempts
COCO_configurations = msCOCO_matrix.get_max_co_occurrence(5, int(objects_in_scene + 5))
configuration_1 = COCO_configurations[0]
print("Config 1:", configuration_1)
# TDW models/objects
objects = []
for COCO_object in configuration_1:
print(COCO_object)
print(COCO_object.split())
if len(COCO_object.split()) > 1:
COCO_object = COCO_object.split()[-1]
print(COCO_object)
# Check if COCO category is a key in the COCO-TDW mapping
if COCO_object in TDW_COCO_models.keys():
# Gets random TDW model (from COCO-to-TDW map) based on COCO category key
print(TDW_COCO_models[COCO_object])
model = TDW_COCO_models[COCO_object][random.randint(0, len(TDW_COCO_models[COCO_object]) - 1)]
objects.append(model)
print("COCO to TDW objects:", objects)
# print(len(objects))
# Stores object categories that other objects can be placed upon (e.g. table, chair, couch, bed)
surface_properties_list = TDW_COCO_models['table'] + TDW_COCO_models['chair'] + \
TDW_COCO_models['bed'] + TDW_COCO_models['couch'] + \
TDW_COCO_models['bench'] + TDW_COCO_models['refrigerator']
surface_categories = []
for surface_properties in surface_properties_list:
surface_categories.append(surface_properties[0])
print("Surface Categories:", surface_categories)
# Stores the actual surface object instances/ids alongside number of objects on the surface
surface_object_ids = {}
self.start()
positions_list = [] # Stores current model locations and radii
scene_dimensions = [] # Store scene/environment dimensions
init_setup_commands = [{"$type": "set_screen_size",
"width": 640,
"height": 481},
{"$type": "set_render_quality",
"render_quality": 1}]
self.communicate(init_setup_commands)
scene_lib = SceneLibrarian()
# Disable physics when adding in new objects (objects float)
self.communicate({"$type": "simulate_physics",
"value": False})
for scene in scenes[1:]:
# Load in scene
# print("Scene", scene[0])
if scene[3] == "interior" and scene[0] == "box_room_2018":
self.start()
scene_record = scene_lib.get_record(scene[0])
self.communicate({"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()})
# Gets dimensions of environments (e.g. inside, outside) in the scene
# This command returns environment data in the form of a list of serialized byte arrays
scene_bytes = self.communicate({"$type": "send_environments",
"frequency": "once"})
# Iterating through data and parsing byte array
# Ignoring the last element (the frame count)
for b in scene_bytes[:-1]:
e = Environments(b)
for i in range(e.get_num()):
center = e.get_center(i)
bounds = e.get_bounds(i)
env_id = e.get_id(i)
scene_dimensions = [center, bounds, env_id] # Center, bounds are tuples
# Must come before set_pass_masks
avatar_position = TDWUtils.array_to_vector3([0.9 * scene_dimensions[1][0] / 2,
scene_dimensions[1][1] / 2,
0])
# print("Avatar Position:", avatar_position)
self.communicate(TDWUtils.create_avatar(avatar_id="avatar",
position=avatar_position,
look_at={"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0}))
# Set collision mode
self.communicate({"$type": "set_avatar_collision_detection_mode",
"mode": "continuous_speculative",
"avatar_id": "avatar"})
# Alter FOV
self.communicate({"$type": "set_field_of_view",
"field_of_view": 80,
"avatar_id": "avatar"})
# # Gets rid of header (Model: Category)
# objects = TDW_COCO_models[1:]
# random.shuffle(objects)
obj_count = 0
obj_overlaps = 0 # Number of failed attempts to place object due to over-dense objects area
while obj_count < objects_in_scene and obj_overlaps < 5:
# Handles if object has been added to a flat surface
added_to_surface = False
print("Object COUNT:", obj_count)
# Need to have random position for Bounds Data to return meaningful info
valid_obj_pos = {"x": random.uniform(-1 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][0] / 2),
"y": scene_dimensions[1][1] / 4,
"z": random.uniform(-0.9 * scene_dimensions[1][2] / 2,
0.9 * scene_dimensions[1][2] / 2)}
print("First random position")
# Add in the object at random position
# Object will later be removed or updated accordingly after performing collision calculations
record = ModelLibrarian(library="models_full.json").get_record(objects[obj_count][0])
print("Record gotten")
print(objects[obj_count][0])
o_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": TDWUtils.VECTOR3_ZERO,
"category": record.wcategory,
"id": obj_count})
print("Random first add")
# Returns bound data for added object
bounds_data = self.communicate({"$type": "send_bounds",
"frequency": "once"})
print("Bounds returned")
# Appends object, with information on position and obj_radius, to positions_list
# Length of buffer array should be 1
# print("Bounds Data:", bounds_data)
for b in bounds_data[:-1]:
# print("Buffer Loop:", b)
b_id = OutputData.get_data_type_id(b)
if b_id == "boun":
# print("BOUNDS")
o = Bounds(b)
# print("# of Objects:", o.get_num())
# print("# of Failed Attempts:", obj_overlaps)
# print("Buffer Array:", b)
# print("Bounds Object:", o)
for i in range(o.get_num()):
print("Object ID:", o.get_id(i))
print("obj_count:", obj_count)
print("Object:", objects[obj_count][0], "Category:", objects[obj_count][1])
# print("Object Center:", o.get_center(i))
# Only want to compute valid_position for object we are about to add
# Skip any computation if this is not the case
if o.get_id(i) != obj_count:
continue
# Useful for detecting if object fits in environment
# print("Calculating if object fits in environment")
width = distance.euclidean(o.get_left(i), o.get_right(i))
depth = distance.euclidean(o.get_front(i), o.get_back(i))
height = distance.euclidean(o.get_top(i), o.get_bottom(i))
# print("Width:", width)
# print("Depth:", depth)
# ("Height:", height)
# Useful for avoiding object overlap
# print("Calculating Object Bounds")
center_to_top = distance.euclidean(o.get_center(i), o.get_top(i))
center_to_bottom = distance.euclidean(o.get_center(i), o.get_bottom(i))
center_to_left = distance.euclidean(o.get_center(i), o.get_left(i))
center_to_right = distance.euclidean(o.get_center(i), o.get_right(i))
center_to_front = distance.euclidean(o.get_center(i), o.get_front(i))
center_to_back = distance.euclidean(o.get_center(i), o.get_back(i))
# Max object radius (center to diagonal of bounding box)
obj_radius = \
max(math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_back ** 2))
# print("Obj_Radius:", obj_radius)
# Set sweeping radius, based on scene plane dimensions
l_radius = random.uniform(0, min(0.5 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][2] / 2))
# Checking that object fits in scene viewing
if (width > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or
depth > min(0.7 * scene_dimensions[1][0], 0.7 * scene_dimensions[1][2]) or
height > 0.7 * scene_dimensions[1][1]):
print("Object does not fit in scene")
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Not possible to find valid object position -- too many overlapping objects
elif (not self._get_object_position(scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[0]):
print("Could not calculate valid object location")
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
obj_overlaps += 1
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Find appropriate, non-overlapping object position
# Reset object position to the valid position
else:
print("Object fits in scene")
# Check if object fits on table, chair, couch, etc.
# Add object if it fits, place it somewhere on top of the surface
for surface_id in surface_object_ids.keys():
print("Surface ID:", surface_id)
# Skip placement feasibility if the object is already a surface-type object
# Ex. no chair on top of a table
if objects[obj_count][0] in surface_categories:
print("Object: %s is already a surface object" % objects[obj_count][0])
break
# Check how many objects are on surface
if surface_object_ids[surface_id] >= 3:
print("Too many objects on surface")
print("From surface objects dict:", surface_object_ids[surface_id])
continue
surface_bounds = self.get_bounds_data(surface_id)
surface_area = distance.euclidean(surface_bounds.get_left(0),
surface_bounds.get_right(0)) * \
distance.euclidean(surface_bounds.get_front(0),
surface_bounds.get_back(0))
obj_area = width * height
if obj_area < surface_area:
s_center_to_top = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_top(0))
s_center_to_bottom = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_bottom(0))
s_center_to_left = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_left(0))
s_center_to_right = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_right(0))
s_center_to_front = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_front(0))
s_center_to_back = distance.euclidean(surface_bounds.get_center(0),
surface_bounds.get_back(0))
surface_radius = \
max(math.sqrt(
s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_top ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_front ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_left ** 2 + s_center_to_back ** 2),
math.sqrt(
s_center_to_bottom ** 2 + s_center_to_right ** 2 + s_center_to_back ** 2))
print("Surface-type object")
self.communicate({"$type": "destroy_object",
"id": obj_count})
# Adding the object to the top of the surface
on_pos = surface_bounds.get_top(0)
on_y = on_pos[1]
on_pos = TDWUtils.get_random_point_in_circle(np.array(on_pos),
0.7 * surface_radius)
on_pos[1] = on_y
on_pos = TDWUtils.array_to_vector3(on_pos)
on_rot = {"x": 0, "y": random.uniform(-45, 45), "z": 0}
# Add the object.
print("Model Name on Surface:", objects[obj_count][0])
record = ModelLibrarian(library="models_full.json").get_record(
objects[obj_count][0])
on_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": on_pos,
"rotation": on_rot,
"category": record.wcategory,
"id": obj_count})
obj_count += 1
surface_object_ids[surface_id] += 1
object_ids.append(obj_count)
print("Object added on top of surface")
added_to_surface = True
# Breaking out of surface objects loop
break
if added_to_surface:
print("Breaking out of object loop")
# Breaking out of object loop
break
print("Post-surface")
valid_obj_pos = self._get_object_position(scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[1]
print("Position calculated")
positions_list.append(ObjectPosition(valid_obj_pos, obj_radius))
self.communicate([{"$type": "send_images",
"frequency": "never"},
{"$type": "destroy_object",
"id": obj_count}])
added_object_id = self.communicate({"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": {"x": 0, "y": 0, "z": 0},
"category": record.wcategory,
"id": obj_count})
# print("Object ID:", added_object_id)
print("Regular object add")
object_ids.append(added_object_id)
# If TDW model belongs to surface categories, store id_information
if objects[obj_count][0] in surface_categories:
surface_object_ids[obj_count] = 0
# Rotate the object randomly
print("Rotating object")
self.communicate({"$type": "rotate_object_by",
"angle": random.uniform(-45, 45),
"axis": "yaw",
"id": obj_count,
"is_world": True})
# Minimal rotating for position differences
# Don't rotate the object if doing so will result in overlap into scene
if not (o.get_bottom(i)[1] < 0 or o.get_top(i)[1] > 0.9 * scene_dimensions[1][1]):
pitch_angle = random.uniform(-45, 45)
self.communicate({"$type": "rotate_object_by",
"angle": pitch_angle,
"axis": "pitch",
"id": obj_count,
"is_world": True})
roll_angle = random.uniform(-45, 45)
self.communicate({"$type": "rotate_object_by",
"angle": roll_angle,
"axis": "roll",
"id": obj_count,
"is_world": True})
# Don't need this for just changing positions
# Setting random materials/textures
# Looping through sub-objects and materials
sub_count = 0
for sub_object in record.substructure:
# Loop through materials in sub-objects
for j in range(len(sub_object)):
# Get random material and load in
material = random.choice(materials[1:])
self.load_material(material)
print("Material loaded")
# Set random material on material of sub-object
self.communicate({"$type": "set_visual_material",
"material_index": j,
"material_name": material[0],
"object_name": sub_object['name'],
"id": obj_count})
print("Material set")
sub_count += 1
if sub_count > 10:
break
break
print("Updating count")
obj_count += 1
print("Breaking out of object_id loop")
break
# Break out of buffer loop
print("Breaking out of buffer loop")
break
# Move onto next iteration of while loop (next object to load in)
print("Object added - next while loop iteration")
continue
# for i in range(200):
# self.communicate({"$type": "simulate_physics",
# "value": False})
# Enable image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
# Capture scene
# NOTE: THESE SCENES GET REPLACED IN THE TARGET DIRECTORY
scene_data = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": {"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0}})
images = Images(scene_data[0])
TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=path)
print("Object ids:", object_ids)
from typing import Dict, List
import random
from tdw.librarian import ModelLibrarian
from tdw.tdw_utils import TDWUtils
# Every model library, sorted by name.
MODEL_LIBRARIES: Dict[str, ModelLibrarian] = {}
for filename in ModelLibrarian.get_library_filenames():
MODEL_LIBRARIES.update({filename: ModelLibrarian(filename)})
def get_move_along_direction(pos: Dict[str, float], target: Dict[str, float], d: float, noise: float = 0) -> \
Dict[str, float]:
"""
:param pos: The object's position.
:param target: The target position.
:param d: The distance to teleport.
:param noise: Add a little noise to the teleport.
:return: A position from pos by distance d along a directional vector defined by pos, target.
"""
direction = TDWUtils.array_to_vector3(
(TDWUtils.vector3_to_array(target) - TDWUtils.vector3_to_array(pos)) /
TDWUtils.get_distance(pos, target))
return {
"x": pos["x"] + direction["x"] * d + random.uniform(-noise, noise),
"y": pos["y"],
"z": pos["z"] + direction["z"] * d + random.uniform(-noise, noise)
}
class Transport(Magnebot):
"""
Transport challenge API.
```python
from transport_challenge import Transport
m = Transport()
# Initializes the scene.
status = m.init_scene(scene="2a", layout=1)
print(status) # ActionStatus.success
# Prints the current position of the Magnebot.
print(m.state.magnebot_transform.position)
# Prints a list of all container IDs.
print(m.containers)
```
**This extends the Magnebot API. Please read the [Magnebot API documentation](https://github.com/alters-mit/magnebot/blob/main/doc/magnebot_controller.md).**
This API includes the following changes and additions:
- Procedurally add **containers** and **target objects** to the scene. Containers are boxes without lids that can hold objects; see the `containers` field. Target objects are small objects that must be transported to the goal zone; see the `target_objects` field. These containers and target objects are included alongside all other objects in [`self.objects_static` and `self.state`](https://github.com/alters-mit/magnebot/blob/main/doc/magnebot_controller.md#fields). - Higher-level actions to pick up target objects and put them in containers.
- A few new actions: `pick_up()`, `put_in()`, and `pour_out()`
- Modified behavior for certain Magnebot actions such as `reset_arm()`
- An interaction budget. The field `action_cost` increments by an action's "cost" at the end of the action:
| Action | Cost |
| --- | --- |
| `init_scene()` | 0 |
| `turn_by()` | 1 |
| `turn_to()` | 1 |
| `move_by()` | 1 |
| `move_to()` | 2 |
| `reset_position()` | 1 |
| `reach_for()` | 1 |
| `grasp()` | 1 |
| `drop()` | 1 |
| `reset_arm()` | 1 |
| `rotate_camera()` | 0 |
| `reset_camera()` | 0 |
| `add_camera()` | 0 |
| `end()` | 0 |
| `pick_up()` | 2 |
| `put_in()` | 1 |
| `pour_out()` | 1 |
| `get_target_objects_in_goal_zone()` | 0 |
"""
""":class_var
The mass of each target object.
"""
TARGET_OBJECT_MASS: float = 0.25
""":class_var
The goal zone is a circle defined by `self.goal_center` and this radius value.
"""
GOAL_ZONE_RADIUS: float = 1
# The scale factor of each container relative to its original size.
__CONTAINER_SCALE = {"x": 0.6, "y": 0.4, "z": 0.6}
# The mass of a container.
__CONTAINER_MASS = 1
# The model librarian.
__LIBRARIAN = ModelLibrarian()
# The value of the torso prismatic joint while the Magnebot is holding a container.
__TORSO_PRISMATIC_CONTAINER = 1.2
# Load a list of visual materials for target objects.
__TARGET_OBJECT_MATERIALS = TARGET_OBJECT_MATERIALS_PATH.read_text(
encoding="utf-8").split("\n")
def __init__(self,
port: int = 1071,
launch_build: bool = False,
screen_width: int = 256,
screen_height: int = 256,
debug: bool = False,
auto_save_images: bool = False,
images_directory: str = "images",
random_seed: int = None,
img_is_png: bool = True,
skip_frames: int = 10):
"""
:param port: The socket port. [Read this](https://github.com/threedworld-mit/tdw/blob/master/Documentation/getting_started.md#command-line-arguments) for more information.
:param launch_build: If True, the build will launch automatically on the default port (1071). If False, you will need to launch the build yourself (for example, from a Docker container).
:param screen_width: The width of the screen in pixels.
:param screen_height: The height of the screen in pixels.
:param auto_save_images: If True, automatically save images to `images_directory` at the end of every action.
:param images_directory: The output directory for images if `auto_save_images == True`.
:param debug: If True, enable debug mode. This controller will output messages to the console, including any warnings or errors sent by the build. It will also create 3D plots of arm articulation IK solutions.
:param random_seed: The random seed used for setting the start position of the Magnebot, the goal room, and the target objects and containers.
:param img_is_png: If True, the `img` pass images will be .png files. If False, the `img` pass images will be .jpg files, which are smaller; the build will run approximately 2% faster.
:param skip_frames: The build will return output data this many frames per `communicate()` call. This will greatly speed up the simulation. If you want to render every frame, set this to 0.
"""
super().__init__(port=port,
launch_build=launch_build,
screen_width=screen_width,
screen_height=screen_height,
debug=debug,
auto_save_images=auto_save_images,
images_directory=images_directory,
random_seed=random_seed,
img_is_png=img_is_png,
skip_frames=skip_frames)
""":field
The IDs of each target object in the scene.
"""
self.target_objects: List[int] = list()
""":field
The IDs of each container in the scene.
"""
self.containers: List[int] = list()
""":field
The total number of actions taken by the Magnebot.
"""
self.action_cost: int = 0
""":field
The challenge is successful when the Magnebot moves all of the target objects to the the goal zone, which is defined by this position and `Transport.GOAL_ZONE_RADIUS`. This value is set in `init_scene()`.
"""
self.goal_position: np.array = np.array([0, 0, 0])
""":field
The room that `self.goal_position` is in. [See here](https://github.com/alters-mit/magnebot/tree/main/doc/images/rooms) for images of the rooms. This value is set in `init_scene()`.
"""
self.goal_room: int = 0
""":field
If True, the Magnebot successfully transported all of the objects to the goal zone. This is updated at the end of every action, including actions with 0 cost.
"""
self.done: bool = False
# Cached IK solution for resetting an arm holding a container.
self._container_arm_reset_angles: Dict[Arm, np.array] = dict()
# Get all possible target objects. Key = name. Value = scale.
self._target_objects: Dict[str, float] = dict()
with open(str(TARGET_OBJECTS_PATH.resolve())) as csvfile:
reader = DictReader(csvfile)
for row in reader:
self._target_objects[row["name"]] = float(row["scale"])
self._target_object_names = list(self._target_objects.keys())
def init_scene(self,
scene: str,
layout: int,
room: int = None,
goal_room: int = None) -> ActionStatus:
"""
This is the same function as `Magnebot.init_scene()` but it adds target objects and containers to the scene.
When `init_scene()` is called, 8-12 target objects will be randomly placed on the floor of a randomly-selected room. Then, there is a 25% chance of adding one container per room.
:param scene: The name of an interior floorplan scene. Each number (1, 2, etc.) has a different shape, different rooms, etc. Each letter (a, b, c) is a cosmetically distinct variant with the same floorplan.
:param layout: The furniture layout of the floorplan. Each number (0, 1, 2) will populate the floorplan with different furniture in different positions.
:param room: The index of the room that the Magnebot will spawn in the center of. If None, the room will be chosen randomly.
:param goal_room: The goal room. If None, this is chosen randomly. See field descriptions of `goal_room` and `goal_position` in this document.
Possible [return values](https://github.com/alters-mit/magnebot/blob/main/doc/action_status.md):
- `success`
:return: An `ActionStatus` (always success).
"""
# Set the room of the goal.
rooms = np.unique(
np.load(
str(ROOM_MAPS_DIRECTORY.joinpath(
f"{scene[0]}.npy").resolve())))
if goal_room is None:
self.goal_room = int(self._rng.choice(rooms))
else:
assert goal_room in rooms, f"Not a valid room: {goal_room}"
self.goal_room = goal_room
# The goal position is the center of the room.
self.goal_position = TDWUtils.vector3_to_array(
loads(SPAWN_POSITIONS_PATH.read_text())[scene[0]][str(layout)][str(
self.goal_room)])
if self._debug:
print(f"Goal position: {self.goal_position}")
return super().init_scene(scene=scene, layout=layout, room=room)
def pick_up(self, target: int, arm: Arm) -> ActionStatus:
"""
Grasp an object and lift it up. This combines the actions `grasp()` and `reset_arm()`.
Possible [return values](https://github.com/alters-mit/magnebot/blob/main/doc/action_status.md):
- `success`
- `cannot_reach`
- `failed_to_grasp` (Either because the motion failed or because the magnet is already holding a different object.)
- `failed_to_bend`
:param target: The ID of the target object.
:param arm: The arm of the magnet that will try to grasp the object.
:return: An `ActionStatus` indicating if the magnet at the end of the `arm` is holding the `target` and if not, why.
"""
if target in self.state.held[arm]:
if self._debug:
print(f"Already holding {target}")
return ActionStatus.success
if len(self.state.held[arm]) > 0:
if self._debug:
print(f"Already holding an object in {arm.name}")
return ActionStatus.failed_to_grasp
# This will increment `self.action_cost`.
grasp_status = self.grasp(target=target, arm=arm)
reset_status = self.reset_arm(arm=arm, reset_torso=True)
self._end_action()
if grasp_status != ActionStatus.success:
return grasp_status
return reset_status
def reset_arm(self, arm: Arm, reset_torso: bool = True) -> ActionStatus:
"""
This is the same as `Magnebot.reset_arm()` unless the arm is holding a container.
If the arm is holding a container, it will try to align the bottom of the container with the floor. This will be somewhat slow the first time the Magnebot does this for this held container.
Possible [return values](https://github.com/alters-mit/magnebot/blob/main/doc/action_status.md):
- `success`
- `failed_to_bend`
:param arm: The arm that will be reset.
:param reset_torso: If True, rotate and slide the torso to its neutral rotation and height.
:return: An `ActionStatus` indicating if the arm reset and if not, why.
"""
self.action_cost += 1
# Use cached angles to reset an arm holding a container.
if arm in self._container_arm_reset_angles:
return super().reset_arm(arm=arm, reset_torso=reset_torso)
status = super().reset_arm(arm=arm, reset_torso=reset_torso)
for object_id in self.state.held[arm]:
# If the arm is holding a container, orient the container to be level with the floor.
if object_id in self.containers:
rot = self.state.object_transforms[object_id].rotation
# Source: https://answers.unity.com/questions/416169/finding-pitchrollyaw-from-quaternions.html
x_rot = -np.rad2deg(
np.arctan2(2 * rot[0] * rot[3] - 2 * rot[1] * rot[2],
1 - 2 * rot[0] * rot[0] - 2 * rot[2] * rot[2]))
if x_rot > 90:
x_rot = x_rot - 180
elif x_rot > 0:
x_rot = -x_rot
if self._debug:
print(f"Setting x rotation of container wrist to {x_rot}")
# Get the commands to reset the arm.
self._next_frame_commands.extend(
self._get_reset_arm_commands(arm=arm,
reset_torso=reset_torso))
# Get the ID of the wrist.
if arm == Arm.right:
wrist_id = self.magnebot_static.arm_joints[
ArmJoint.wrist_right]
else:
wrist_id = self.magnebot_static.arm_joints[
ArmJoint.wrist_left]
temp = list()
for cmd in self._next_frame_commands:
# Adjust the wrist to level off the container.
if cmd["$type"] == "set_spherical_target" and cmd[
"joint_id"] == wrist_id:
cmd["target"] = {"x": x_rot, "y": 0, "z": 0}
temp.append(cmd)
self._next_frame_commands = temp
self._start_action()
# Bend the arm.
self._do_arm_motion(joint_ids=[wrist_id])
self._end_action()
# Cache the arm angles so we can next time immediately reset to this position.
self._container_arm_reset_angles[arm] = np.array([
np.rad2deg(a)
for a in self._get_initial_angles(arm=arm)[1:-1]
])
return status
return status
def put_in(self) -> ActionStatus:
"""
Put an object in a container. In order to put an object in a container:
- The Magnebot must be holding a container with one magnet.
- The Magnebot must be holding a target object with another magnet.
This is a multistep action, combining many motions and may require more time than other actions.
Possible [return values](https://github.com/alters-mit/magnebot/blob/main/doc/action_status.md):
- `success`
- `not_holding` (If the Magnebot isn't holding a container or target object.)
- `not_in` (If the target object didn't land in the container.)
:return: An `ActionStatus` indicating if the target object is in the container and if not, why.
"""
def __object_in_container(s: SceneState) -> bool:
"""
:param s: The scene state.
:return: True if the target object is on the surface of the container or collides with an object in the container.
"""
# The ID of the target object is actually the ID of the magnet because it's attached to the magnet.
magnet_id = self.magnebot_static.magnets[object_arm]
# Is the target object colliding with the surface of the container?
return container_id in self._trigger_events and (
magnet_id in self._trigger_events[container_id]
or object_id in self._trigger_events[container_id])
# Get the arm holding each object.
container_arm, container_id = self._get_container_arm()
if container_arm is None:
if self._debug:
print("Magnebot isn't holding a container.")
return ActionStatus.not_holding
# Check whether the opposite arm is holding a target object.
object_arm = Arm.left if container_arm == Arm.right else Arm.right
object_id = None
for o_id in self.state.held[object_arm]:
if o_id in self.target_objects:
object_id = o_id
if object_id is None:
if self._debug:
print(
f"Magnebot is holding a container with the {container_arm.name} magnet but isn't holding a "
f"target object with the {object_arm.name} magnet.")
return ActionStatus.not_holding
self._start_action()
self._next_frame_commands.append({
"$type": "set_immovable",
"immovable": True
})
# Move the other arm out of the way.
if container_arm == Arm.right:
elbow_id = self.magnebot_static.arm_joints[ArmJoint.elbow_left]
else:
elbow_id = self.magnebot_static.arm_joints[ArmJoint.elbow_right]
self._next_frame_commands.append({
"$type": "set_revolute_target",
"joint_id": elbow_id,
"target": 115
})
state = SceneState(resp=self.communicate([]))
# Bring the container approximately to center.
ct = {
"x": 0.1 * (-1 if container_arm is Arm.right else 1),
"y": 0.4,
"z": 0.5
}
self._start_ik(
target=ct,
arm=container_arm,
absolute=False,
allow_column=False,
state=state,
fixed_torso_prismatic=Transport.__TORSO_PRISMATIC_CONTAINER,
do_prismatic_first=False)
self._do_arm_motion()
state = SceneState(resp=self.communicate([]))
# Move the target object to be over the container.
target = np.copy(state.object_transforms[container_id].position)
target[1] += 0.5
self._start_ik(
target=TDWUtils.array_to_vector3(target),
arm=object_arm,
allow_column=False,
state=state,
absolute=True,
fixed_torso_prismatic=Transport.__TORSO_PRISMATIC_CONTAINER,
object_id=object_id)
# Get the ID of the wrist.
if object_arm == Arm.right:
wrist_id = self.magnebot_static.arm_joints[ArmJoint.wrist_right]
else:
wrist_id = self.magnebot_static.arm_joints[ArmJoint.wrist_left]
# Move the wrist down for a better angle.
self._next_frame_commands.extend([{
"$type": "set_spherical_target",
"joint_id": wrist_id,
"target": {
"x": -45,
"y": 0,
"z": 0
}
}])
self._do_arm_motion(conditional=__object_in_container,
joint_ids=[wrist_id])
# Drop the object.
self._append_drop_commands(object_id=object_id, arm=object_arm)
# Set the detection mode to discrete. This will make physics less buggy.
self._next_frame_commands.append({
"$type": "set_object_collision_detection_mode",
"id": int(object_id),
"mode": "discrete"
})
# Wait for the object to fall (hopefully into the container).
self._wait_until_objects_stop(object_ids=[object_id],
state=SceneState(self.communicate([])))
# Reset the arms.
self.reset_arm(arm=object_arm, reset_torso=False)
self.reset_arm(arm=container_arm, reset_torso=True)
self.action_cost -= 1
in_container = self._get_objects_in_container(
container_id=container_id)
# If the object isn't in in the container, set the detection mode to the default.
if object_id not in in_container:
self._next_frame_commands.append({
"$type": "set_object_collision_detection_mode",
"id": int(object_id),
"mode": "continuous_dynamic"
})
self._end_action()
if object_id in in_container:
return ActionStatus.success
else:
if self._debug:
print(f"Object {object_id} isn't in container {container_id}")
return ActionStatus.not_in
def pour_out(self) -> ActionStatus:
"""
Pour out all of the objects in a container held by one of the Magnebot's magnets.
The Magnebot will extend the arm holding the container and then flip its elbow and wrist.
The action ends when any objects that were in the container stop moving.
Possible [return values](https://github.com/alters-mit/magnebot/blob/main/doc/action_status.md):
- `success`
- `not_holding` (If the Magnebot isn't holding a container.)
- `still_in` (If there are objects still in the container.)
:return: An `ActionStatus` indicating whether the container is now empty and if not, why.
"""
container_arm, container_id = self._get_container_arm()
if container_arm is None:
if self._debug:
print("Magnebot isn't holding a container.")
return ActionStatus.not_holding
# Get all of the objects currently in the container.
in_container_0 = self._get_objects_in_container(
container_id=container_id)
self._start_action()
self._next_frame_commands.append({
"$type": "set_immovable",
"immovable": True
})
# Get the joint IDs.
if container_arm == Arm.right:
shoulder_id = self.magnebot_static.arm_joints[
ArmJoint.shoulder_right]
elbow_id = self.magnebot_static.arm_joints[ArmJoint.elbow_right]
wrist_id = self.magnebot_static.arm_joints[ArmJoint.wrist_right]
else:
shoulder_id = self.magnebot_static.arm_joints[
ArmJoint.shoulder_left]
elbow_id = self.magnebot_static.arm_joints[ArmJoint.elbow_left]
wrist_id = self.magnebot_static.arm_joints[ArmJoint.wrist_left]
# Extend the arm.
self._next_frame_commands.extend([{
"$type": "set_spherical_target",
"joint_id": shoulder_id,
"target": {
"x": -90,
"y": 0,
"z": 0
}
}, {
"$type": "set_revolute_target",
"joint_id": elbow_id,
"target": 0
}])
self._do_arm_motion(joint_ids=[shoulder_id, elbow_id])
# Flip the wrist and elbow.
self._next_frame_commands.extend([{
"$type": "set_spherical_target",
"joint_id": wrist_id,
"target": {
"x": 90,
"y": 0,
"z": 0
}
}, {
"$type": "set_revolute_target",
"joint_id": elbow_id,
"target": 35
}])
self._do_arm_motion(joint_ids=[wrist_id, elbow_id])
# Wait for the objects to fall out (by this point, they likely already have).
self._wait_until_objects_stop(in_container_0,
state=SceneState(self.communicate([])))
self._next_frame_commands.extend(
self._get_reset_arm_commands(arm=container_arm, reset_torso=False))
self._do_arm_motion()
in_container_1 = self._get_objects_in_container(
container_id=container_id)
# Reset the collision detection mode of objects that were poured out.
for object_id in in_container_0:
if object_id not in in_container_1:
self._next_frame_commands.append({
"$type": "set_object_collision_detection_mode",
"id": int(object_id),
"mode": "continuous_dynamic"
})
self._end_action()
self.action_cost += 1
if len(in_container_1) == 0:
return ActionStatus.success
else:
return ActionStatus.still_in
def get_target_objects_in_goal_zone(self) -> List[int]:
"""
:return: A list of IDs of all of the target objects currently in the goal zone.
"""
objects: List[int] = list()
# Objects that are still being held by the Magnebot don't count.
held: List[int] = list()
for arm in self.state.held:
for object_id in self.state.held[arm]:
if object_id in self.target_objects:
held.append(object_id)
# The object must be in the goal zone and on the floor.
for object_id in self.target_objects:
if self.state.object_transforms[object_id].position[1] <= 0.1 and \
np.linalg.norm(self.state.object_transforms[object_id].position - self.goal_position) <= \
Transport.GOAL_ZONE_RADIUS:
objects.append(object_id)
return objects
def drop(self,
target: int,
arm: Arm,
wait_for_objects: bool = True) -> ActionStatus:
status = super().drop(target=target,
arm=arm,
wait_for_objects=wait_for_objects)
if status == ActionStatus.success:
# Remove the cached container arm angles.
if arm in self._container_arm_reset_angles:
del self._container_arm_reset_angles[arm]
self.action_cost += 1
return status
def turn_by(self, angle: float, aligned_at: float = 3) -> ActionStatus:
self.action_cost += 1
return super().turn_by(angle=angle, aligned_at=aligned_at)
def turn_to(self,
target: Union[int, Dict[str, float]],
aligned_at: float = 3) -> ActionStatus:
self.action_cost += 1
return super().turn_to(target=target, aligned_at=aligned_at)
def move_by(self,
distance: float,
arrived_at: float = 0.3) -> ActionStatus:
self.action_cost += 1
return super().move_by(distance=distance, arrived_at=arrived_at)
def reach_for(self,
target: Dict[str, float],
arm: Arm,
absolute: bool = True,
arrived_at: float = 0.125) -> ActionStatus:
self.action_cost += 1
return super().reach_for(target=target,
arm=arm,
absolute=absolute,
arrived_at=arrived_at)
def grasp(self, target: int, arm: Arm) -> ActionStatus:
self.action_cost += 1
return super().grasp(target=target, arm=arm)
def reset_position(self) -> ActionStatus:
self.action_cost += 1
return super().reset_position()
def get_scene_init_commands(self, scene: str, layout: int,
audio: bool) -> List[dict]:
# Clear the list of target objects and containers.
self.target_objects.clear()
self.containers.clear()
commands = super().get_scene_init_commands(scene=scene,
layout=layout,
audio=audio)
# Load the map of the rooms in the scene, the occupancy map, and the scene bounds.
room_map = np.load(
str(ROOM_MAPS_DIRECTORY.joinpath(f"{scene[0]}.npy").resolve()))
self.occupancy_map = np.load(
str(
OCCUPANCY_MAPS_DIRECTORY.joinpath(
f"{scene[0]}_{layout}.npy").resolve()))
self._scene_bounds = loads(SCENE_BOUNDS_PATH.read_text())[scene[0]]
# Sort all free positions on the occupancy map by room.
rooms: Dict[int, List[Tuple[int, int]]] = dict()
for ix, iy in np.ndindex(room_map.shape):
room_index = room_map[ix][iy]
if room_index not in rooms:
rooms[room_index] = list()
if self.occupancy_map[ix][iy] == 0:
rooms[room_index].append((ix, iy))
# Choose a random room.
target_room_index = self._rng.choice(np.array(list(rooms.keys())))
target_room_positions: np.array = np.array(rooms[target_room_index])
used_target_object_positions: List[Tuple[int, int]] = list()
# Add target objects to the room.
for i in range(self._rng.randint(8, 12)):
got_position = False
ix, iy = -1, -1
# Get a position where there isn't a target object.
while not got_position:
ix, iy = target_room_positions[self._rng.randint(
0, len(target_room_positions))]
got_position = True
for utop in used_target_object_positions:
if utop[0] == ix and utop[1] == iy:
got_position = False
used_target_object_positions.append((ix, iy))
# Get the (x, z) coordinates for this position.
x, z = self.get_occupancy_position(ix, iy)
self._add_target_object(model_name=self._rng.choice(
self._target_object_names),
position={
"x": x,
"y": 0,
"z": z
})
# Add containers throughout the scene.
containers = CONTAINERS_PATH.read_text(encoding="utf-8").split("\n")
for room_index in list(rooms.keys()):
# Maybe don't add a container in this room.
if self._rng.random() < 0.25:
continue
# Get a random position in the room.
room_positions: np.array = np.array(rooms[room_index])
got_position = False
ix, iy = -1, -1
# Get a position where there isn't a target object.
while not got_position:
ix, iy = room_positions[self._rng.randint(
0, len(room_positions))]
got_position = True
for utop in used_target_object_positions:
if utop[0] == ix and utop[1] == iy:
got_position = False
# Get the (x, z) coordinates for this position.
x, z = self.get_occupancy_position(ix, iy)
container_name = self._rng.choice(containers)
self._add_container(model_name=container_name,
position={
"x": x,
"y": 0,
"z": z
},
rotation={
"x": 0,
"y": self._rng.uniform(-179, 179),
"z": 0
})
return commands
def _cache_static_data(self, resp: List[bytes]) -> None:
# Reset the action counter and challenge status.
self.action_cost = 0
self.done = False
super()._cache_static_data(resp=resp)
def _add_container(self,
model_name: str,
position: Dict[str, float] = None,
rotation: Dict[str, float] = None) -> int:
"""
Add a container. Cache the ID.
:param model_name: The name of the container.
:param position: The initial position of the container.
:param rotation: The initial rotation of the container.
:return: The ID of the container.
"""
object_id = self._add_object(position=position,
rotation=rotation,
scale=Transport.__CONTAINER_SCALE,
audio=self._OBJECT_AUDIO[model_name],
model_name=model_name)
self.containers.append(object_id)
# Set a light mass for each container.
self._object_init_commands[object_id].append({
"$type":
"set_mass",
"id":
object_id,
"mass":
Transport.__CONTAINER_MASS
})
# Add a trigger collider inside the container.
# We'll use this trigger collider to determine if an object is in the container..
self._object_init_commands[object_id].extend([{
"$type": "add_trigger_collider",
"id": object_id,
"shape": "cube",
"enter": True,
"stay": True,
"exit": False,
"position": {
"x": 0,
"y": 0.1525,
"z": 0
},
"scale": {
"x": 0.457,
"y": 0.305,
"z": 0.457
}
}])
return object_id
def _add_target_object(self, model_name: str,
position: Dict[str, float]) -> int:
"""
Add a targt object. Cache the ID.
:param model_name: The name of the target object.
:param position: The initial position of the target object.
:return: The ID of the target object.
"""
# Set custom object info for the target objects.
audio = ObjectInfo(name=model_name,
mass=Transport.TARGET_OBJECT_MASS,
material=AudioMaterial.ceramic,
resonance=0.6,
amp=0.01,
library="models_core.json",
bounciness=0.5)
scale = self._target_objects[model_name]
# Add the object.
object_id = self._add_object(position=position,
rotation={
"x": 0,
"y": self._rng.uniform(-179, 179),
"z": 0
},
scale={
"x": scale,
"y": scale,
"z": scale
},
audio=audio,
model_name=model_name)
self.target_objects.append(object_id)
# Set a random visual material for each target object.
visual_material = self._rng.choice(Transport.__TARGET_OBJECT_MATERIALS)
substructure = Transport.__LIBRARIAN.get_record(
model_name).substructure
self._object_init_commands[object_id].extend(
TDWUtils.set_visual_material(substructure=substructure,
material=visual_material,
object_id=object_id,
c=self,
quality="low"))
return object_id
def _get_reset_arm_commands(self, arm: Arm,
reset_torso: bool) -> List[dict]:
if arm in self._container_arm_reset_angles:
self._append_ik_commands(
angles=self._container_arm_reset_angles[arm], arm=arm)
return list()
else:
return super()._get_reset_arm_commands(arm=arm,
reset_torso=reset_torso)
def _get_container_arm(self) -> Tuple[Arm, int]:
"""
:return: Tuple: The arm holding a container, if any; the container ID.
"""
container_arm: Optional[Arm] = None
container_id = -1
# Get an arm holding a container.
for arm in self.state.held:
for o_id in self.state.held[arm]:
if container_arm is None and o_id in self.containers:
container_id = o_id
container_arm = arm
return container_arm, container_id
def _get_objects_in_container(self, container_id: int) -> List[int]:
"""
:param container_id: The ID of the container.
:return: A list of objects in the container.
"""
if container_id not in self._trigger_events:
return list()
else:
return self._trigger_events[container_id]
def _is_challenge_done(self) -> bool:
"""
:return: True if all of the objects have been transported to the goal zone.
"""
return len(self.get_target_objects_in_goal_zone()) == len(
self.target_objects)
def _end_action(self) -> None:
super()._end_action()
self.done = self._is_challenge_done()
def _get_bounds_sides(self,
target: int) -> Tuple[List[np.array], List[bytes]]:
sides, resp = super()._get_bounds_sides(target=target)
# Set the y value to the highest point.
max_y = -np.inf
for s in sides:
if s[1] > max_y:
max_y = s[1]
sides = [np.array((s[0], max_y, s[2])) for s in sides]
# Don't try to pick up the top or bottom of a container.
if target in self.containers:
sides = sides[:-2]
return sides, resp
class TransformInitData:
"""
Basic initialization parameters for an object. Can be converted to and from a list of commands.
This is similar to [`Controller.get_add_object()`](controller.md) except that it includes more parameters.
"""
LIBRARIES: Dict[str, ModelLibrarian] = dict()
for _lib_file in ModelLibrarian.get_library_filenames():
LIBRARIES[_lib_file] = ModelLibrarian(_lib_file)
def __init__(self,
name: str,
library: str = "models_core.json",
scale_factor: Dict[str, float] = None,
position: Dict[str, float] = None,
rotation: Dict[str, float] = None,
kinematic: bool = False,
gravity: bool = True):
"""
:param name: The name of the model.
:param library: The filename of the library containing the model's record.
:param scale_factor: The [scale factor](../api/command_api.md#scale_object).
:param position: The initial position. If None, defaults to: `{"x": 0, "y": 0, "z": 0`}.
:param rotation: The initial rotation as Euler angles or a quaternion. If None, defaults to: `{"w": 1, "x": 0, "y": 0, "z": 0}`
:param kinematic: If True, the object will be [kinematic](../api/command_api.md#set_kinematic_state).
:param gravity: If True, the object won't respond to [gravity](../api/command_api.md#set_kinematic_state).
"""
if position is None:
self.position = TDWUtils.VECTOR3_ZERO
else:
self.position = position
if rotation is None:
self.rotation = {"w": 1, "x": 0, "y": 0, "z": 0}
else:
self.rotation = rotation
if scale_factor is None:
self.scale_factor = {"x": 1, "y": 1, "z": 1}
else:
self.scale_factor = scale_factor
self.name = name
self.library = library
self.kinematic = kinematic
self.gravity = gravity
def get_commands(self) -> Tuple[int, List[dict]]:
"""
:return: Tuple: The ID of the object; a list of commands to create the object: `[add_object, rotate_object_to, scale_object, set_kinematic_state, set_object_collision_detection_mode]`
"""
record = TransformInitData.LIBRARIES[self.library].get_record(
name=self.name)
object_id = Controller.get_unique_id()
add_object_command = {
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": self.position,
"category": record.wcategory,
"id": object_id
}
# The rotation is in Euler angles.
if "w" not in self.rotation:
add_object_command["rotation"] = self.rotation
commands = [add_object_command]
# The rotation is a quaternion.
if "w" in self.rotation:
commands.append({
"$type": "rotate_object_to",
"rotation": self.rotation,
"id": object_id
})
commands.extend([{
"$type": "scale_object",
"scale_factor": self.scale_factor,
"id": object_id
}, {
"$type": "set_kinematic_state",
"id": object_id,
"is_kinematic": self.kinematic,
"use_gravity": self.gravity
}])
# Kinematic objects must be continuous_speculative.
if self.kinematic:
commands.append({
"$type": "set_object_collision_detection_mode",
"id": object_id,
"mode": "continuous_speculative"
})
return object_id, commands
def run(self):
rng = np.random.RandomState(0)
self.model_librarian = ModelLibrarian("models_special.json")
self.start()
commands = [TDWUtils.create_empty_room(100, 100)]
# The starting height of the objects.
y = 10
# The radius of the circle of objects.
r = 7.0
# The mass of each object.
mass = 5
# Get all points within the circle defined by the radius.
p0 = np.array((0, 0))
o_id = 0
for x in np.arange(-r, r, 1):
for z in np.arange(-r, r, 1):
p1 = np.array((x, z))
dist = np.linalg.norm(p0 - p1)
if dist < r:
# Add an object.
# Set its mass, physics properties, and color.
commands.extend([self.get_add_object("prim_cone",
object_id=o_id,
position={"x": x, "y": y, "z": z},
rotation={"x": 0, "y": 0, "z": 180}),
{"$type": "set_mass",
"id": o_id,
"mass": mass},
{"$type": "set_physic_material",
"dynamic_friction": 0.8,
"static_friction": 0.7,
"bounciness": 0.5,
"id": o_id},
{"$type": "set_color",
"color": {"r": rng.random_sample(),
"g": rng.random_sample(),
"b": rng.random_sample(),
"a": 1.0},
"id": o_id}])
o_id += 1
# Request transforms per frame.
commands.extend([{"$type": "send_transforms",
"frequency": "always"}])
# Create an avatar to observe the grisly spectacle.
avatar_position = {"x": -20, "y": 8, "z": 18}
commands.extend(TDWUtils.create_avatar(position=avatar_position, look_at=TDWUtils.VECTOR3_ZERO))
commands.append({"$type": "set_focus_distance",
"focus_distance": TDWUtils.get_distance(avatar_position, TDWUtils.VECTOR3_ZERO)})
resp = self.communicate(commands)
# If an objects are this far away from (0, 0, 0) the forcefield "activates".
forcefield_radius = 5
# The forcefield will bounce objects away at this force.
forcefield_force = -10
zeros = np.array((0, 0, 0))
for i in range(1000):
transforms = Transforms(resp[0])
commands = []
for j in range(transforms.get_num()):
pos = transforms.get_position(j)
pos = np.array(pos)
# If the object is in the forcefield, apply a force.
if TDWUtils.get_distance(TDWUtils.array_to_vector3(pos), TDWUtils.VECTOR3_ZERO) <= forcefield_radius:
# Get the normalized directional vector and multiply it by the force magnitude.
d = zeros - pos
d = d / np.linalg.norm(d)
d = d * forcefield_force
commands.append({"$type": "apply_force_to_object",
"id": transforms.get_id(j),
"force": TDWUtils.array_to_vector3(d)})
resp = self.communicate(commands)
def __init__(self, port: int = 1071, launch_build: bool = True):
super().__init__(port=port, launch_build=launch_build)
self.cube = ModelLibrarian("models_flex.json").get_record("cube")
def run(self):
self.start()
positions_list = [] # Stores current model locations and radii
scene_dimensions = [] # Store scene/environment dimensions
init_setup_commands = [{
"$type": "set_screen_size",
"width": 1280,
"height": 962
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
self.communicate(init_setup_commands)
scene_lib = SceneLibrarian()
# Disable physics when adding in new objects (objects float)
self.communicate({"$type": "simulate_physics", "value": False})
for scene in scenes[1:]:
# Load in scene
print("Scene", scene[0])
if scene[3] == "interior" and scene[0] == "box_room_2018":
self.start()
scene_record = scene_lib.get_record(scene[0])
self.communicate({
"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()
})
# Gets dimensions of environments (e.g. inside, outside) in the scene
# This command returns environment data in the form of a list of serialized byte arrays
scene_bytes = self.communicate({
"$type": "send_environments",
"frequency": "once"
})
# Iterating through data and parsing byte array
# Ignoring the last element (the frame count)
for b in scene_bytes[:-1]:
e = Environments(b)
for i in range(e.get_num()):
center = e.get_center(i)
bounds = e.get_bounds(i)
env_id = e.get_id(i)
scene_dimensions = [center, bounds,
env_id] # Center, bounds are tuples
# Must come before set_pass_masks
avatar_position = TDWUtils.array_to_vector3([
0.9 * scene_dimensions[1][0] / 2,
scene_dimensions[1][1] / 2, 0
])
print("Avatar Position:", avatar_position)
self.communicate(
TDWUtils.create_avatar(avatar_id="avatar",
position=avatar_position,
look_at={
"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0
}))
# Set collision mode
self.communicate({
"$type": "set_avatar_collision_detection_mode",
"mode": "continuous_speculative",
"avatar_id": "avatar"
})
# Alter FOV
self.communicate({
"$type": "set_field_of_view",
"field_of_view": 80,
"avatar_id": "avatar"
})
# Gets rid of header (Model: Category)
objects = models[1:]
random.shuffle(objects)
obj_count = 0
obj_overlaps = 0 # Number of failed attempts to place object due to over-dense objects area
while obj_count < 30 and obj_overlaps < 5:
# Need to have random position for Bounds Data to return meaningful info
valid_obj_pos = {
"x":
random.uniform(-1 * scene_dimensions[1][0] / 2,
0.5 * scene_dimensions[1][0] / 2),
"y":
scene_dimensions[1][1] / 4,
"z":
random.uniform(-0.9 * scene_dimensions[1][2] / 2,
0.9 * scene_dimensions[1][2] / 2)
}
# Add in the object at random position
# Object will later be removed or updated accordingly after performing collision calculations
record = ModelLibrarian(
library="models_full.json").get_record(
objects[obj_count][0])
self.communicate({
"$type": "add_object",
"name": objects[obj_count][0],
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": valid_obj_pos,
"rotation": {
"x": 0,
"y": 0,
"z": 0
},
"category": record.wcategory,
"id": obj_count
})
# Returns bound data for added object
bounds_data = self.communicate({
"$type": "send_bounds",
"frequency": "once"
})
# Appends object, with information on position and obj_radius, to positions_list
# Length of buffer array should be 1
print("Bounds Data:", bounds_data)
for b in bounds_data[:-1]:
print("Buffer Loop:", b)
b_id = OutputData.get_data_type_id(b)
if b_id == "boun":
print("BOUNDS")
o = Bounds(b)
print("# of Objects:", o.get_num())
print("# of Failed Attempts:", obj_overlaps)
print("Buffer Array:", b)
print("Bounds Object:", o)
for i in range(o.get_num()):
print("Object ID:", o.get_id(i))
print("obj_count:", obj_count)
print("Object:", objects[obj_count][0],
"Category:", objects[obj_count][1])
print("Object Center:", o.get_center(i))
# Only want to compute valid_position for object we are about to add
# Skip any computation if this is not the case
if o.get_id(i) != obj_count:
continue
# Useful for detecting if object fits in environment
print(
"Calculating if object fits in environment"
)
width = distance.euclidean(
o.get_left(i), o.get_right(i))
depth = distance.euclidean(
o.get_front(i), o.get_back(i))
height = distance.euclidean(
o.get_top(i), o.get_bottom(i))
print("Width:", width)
print("Depth:", depth)
print("Height:", height)
# Useful for avoiding object overlap
print("Calculating Object Bounds")
center_to_top = distance.euclidean(
o.get_center(i), o.get_top(i))
center_to_bottom = distance.euclidean(
o.get_center(i), o.get_bottom(i))
center_to_left = distance.euclidean(
o.get_center(i), o.get_left(i))
center_to_right = distance.euclidean(
o.get_center(i), o.get_right(i))
center_to_front = distance.euclidean(
o.get_center(i), o.get_front(i))
center_to_back = distance.euclidean(
o.get_center(i), o.get_back(i))
# Max object radius (center to diagonal of bounding box)
obj_radius = \
max(math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_top ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_top ** 2 + center_to_right ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_front ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_left ** 2 + center_to_back ** 2),
math.sqrt(center_to_bottom ** 2 + center_to_right ** 2 + center_to_back ** 2))
print("Obj_Radius:", obj_radius)
# Set sweeping radius, based on scene plane dimensions
l_radius = random.uniform(
0,
min(0.9 * scene_dimensions[1][0] / 2,
0.9 * scene_dimensions[1][2] / 2))
# Checking that object fits in scene viewing
if (width > min(0.7 * scene_dimensions[1][0],
0.7 * scene_dimensions[1][2])
or depth > min(
0.7 * scene_dimensions[1][0],
0.7 * scene_dimensions[1][2]) or
height > 0.7 * scene_dimensions[1][1]):
print("Object does not fit in scene")
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Not possible to find valid object position -- too many overlapping objects
elif (not self._get_object_position(
scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[0]):
print(
"Could not calculate valid object location"
)
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
obj_overlaps += 1
# Ensures next attempt to load in item is not the same item as before
random.shuffle(objects)
break
# Find appropriate, non-overlapping object position
# Reset object position to the valid position
else:
print("Object fits in scene")
valid_obj_pos = self._get_object_position(
scene_dimensions=scene_dimensions,
object_positions=positions_list,
object_to_add_radius=obj_radius,
max_tries=20,
location_radius=l_radius)[1]
print("Position calculated")
positions_list.append(
ObjectPosition(valid_obj_pos,
obj_radius))
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": obj_count
}])
print("Object ready to reset")
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type":
"add_object",
"name":
objects[obj_count][0],
"url":
record.get_url(),
"scale_factor":
record.scale_factor,
"position":
valid_obj_pos,
"rotation": {
"x": 0,
"y": 0,
"z": 0
},
"category":
record.wcategory,
"id":
obj_count
}])
print("Object reset")
# Rotate the object randomly
print("Rotating object")
self.communicate({
"$type":
"rotate_object_by",
"angle":
random.uniform(-45, 45),
"axis":
"yaw",
"id":
obj_count,
"is_world":
True
})
# Don't rotate the object if doing so will result in overlap into scene
if not (o.get_bottom(i)[1] < 0
or o.get_top(i)[1] >
0.9 * scene_dimensions[1][1]):
pitch_angle = random.uniform(-45, 45)
self.communicate({
"$type": "rotate_object_by",
"angle": pitch_angle,
"axis": "pitch",
"id": obj_count,
"is_world": True
})
roll_angle = random.uniform(-45, 45)
self.communicate({
"$type": "rotate_object_by",
"angle": roll_angle,
"axis": "roll",
"id": obj_count,
"is_world": True
})
# Setting random materials/textures
# Looping through sub-objects and materials
sub_count = 0
for sub_object in record.substructure:
# Loop through materials in sub-objects
for j in range(len(sub_object)):
# Get random material and load in
material = random.choice(
materials[1:])
self.load_material(material)
print("Material loaded")
# Set random material on material of sub-object
self.communicate({
"$type":
"set_visual_material",
"material_index":
j,
"material_name":
material[0],
"object_name":
sub_object['name'],
"id":
obj_count
})
print("Material set")
sub_count += 1
if sub_count > 10:
break
break
print("Updating count")
obj_count += 1
print("Breaking out of object_id loop")
break
# Break out of buffer loop
print("Breaking out of buffer loop")
break
# Move onto next iteration of while loop (next object to load in)
print("Object added - next while loop iteration")
continue
# Enable image capture
self.communicate({
"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]
})
self.communicate({
"$type": "send_images",
"frequency": "always"
})
# Capture scene
scene_data = self.communicate({
"$type": "look_at_position",
"avatar_id": "avatar",
"position": {
"x": 0,
"y": scene_dimensions[0][1] / 2,
"z": 0
}
})
images = Images(scene_data[0])
TDWUtils.save_images(images,
TDWUtils.zero_padding(i),
output_directory=path)
def __init__(self,
output_dir: Path = Path("D:/audio_dataset"),
total: int = 28602,
port: int = 1071):
"""
:param output_dir: The output directory for the files.
:param port: The socket port.
:param total: The total number of files per sub-set.
"""
self.total = total
self.output_dir = output_dir
if not self.output_dir.exists():
self.output_dir.mkdir(parents=True)
db_path = self.output_dir.joinpath('results.db')
self.conn = sqlite3.connect(str(db_path.resolve()))
self.db_c = self.conn.cursor()
# Sound20K table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sound20k'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE sound20k (path text, scene integer, cam_x real, cam_y real, cam_z real,"
"obj_x real, obj_y real, obj_z real, mass real, static_friction real, dynamic_friction "
"real, yaw real, pitch real, roll real, force real)")
# Scenes table.
if self.db_c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='scenes'").\
fetchone() is None:
self.db_c.execute(
"CREATE TABLE scenes (id integer, commands text)")
self.py_impact = PyImpact()
self.object_info = PyImpact.get_object_info()
sound20k_object_info = PyImpact.get_object_info(
Path("models/object_info.csv"))
for obj_info in sound20k_object_info:
if obj_info in self.object_info:
continue
else:
self.object_info.update(
{obj_info: sound20k_object_info[obj_info]})
self.libs: Dict[str, ModelLibrarian] = {}
# Load all model libraries into memory.
for lib_name in ModelLibrarian.get_library_filenames():
self.libs.update({lib_name: ModelLibrarian(lib_name)})
# Add the custom model library.
self.libs.update({
"models/models.json":
ModelLibrarian(str(Path("models/models.json").resolve()))
})
super().__init__(port=port)
# Global settings.
self.communicate([{
"$type": "set_screen_size",
"width": 256,
"height": 256
}, {
"$type": "set_time_step",
"time_step": 0.02
}, {
"$type": "set_target_framerate",
"framerate": 60
}, {
"$type": "set_physics_solver_iterations",
"iterations": 20
}])
class FlexDominoes(Dominoes, FlexDataset):
FLEX_RECORDS = ModelLibrarian(str(Path("flex.json").resolve())).records
CLOTH_RECORD = MODEL_LIBRARIES["models_special.json"].get_record("cloth_square")
SOFT_RECORD = MODEL_LIBRARIES["models_flex.json"].get_record("sphere")
RECEPTACLE_RECORD = MODEL_LIBRARIES["models_special.json"].get_record("fluid_receptacle1x1")
FLUID_TYPES = FluidTypes()
def __init__(self, port: int = 1071,
all_flex_objects=True,
use_cloth=False,
use_squishy=False,
use_fluid=False,
step_physics=False,
middle_scale_range=0.5,
**kwargs):
Dominoes.__init__(self, port=port, **kwargs)
self._clear_flex_data()
self.all_flex_objects = all_flex_objects
self._set_add_physics_object()
self.step_physics = step_physics
self.use_cloth = use_cloth
self.use_squishy = use_squishy
self.use_fluid = use_fluid
self.middle_scale_range = middle_scale_range
print("MIDDLE SCALE RANGE", self.middle_scale_range)
if self.use_fluid:
self.ft_selection = random.choice(self.FLUID_TYPES.fluid_type_names)
def _set_add_physics_object(self):
if self.all_flex_objects:
self.add_physics_object = self.add_flex_solid_object
self.add_primitive = self.add_flex_solid_object
else:
self.add_physics_object = self.add_rigid_physics_object
def get_scene_initialization_commands(self) -> List[dict]:
commands = Dominoes.get_scene_initialization_commands(self)
commands[0].update({'convexify': True})
create_container = {
"$type": "create_flex_container",
# "collision_distance": 0.001,
"collision_distance": 0.025,
# "collision_distance": 0.1,
"static_friction": 1.0,
"dynamic_friction": 1.0,
"radius": 0.1875,
'max_particles': 50000}
# 'max_particles': 250000}
if self.use_fluid:
create_container.update({
'viscosity': self.FLUID_TYPES.fluid_types[self.ft_selection].viscosity,
'adhesion': self.FLUID_TYPES.fluid_types[self.ft_selection].adhesion,
'cohesion': self.FLUID_TYPES.fluid_types[self.ft_selection].cohesion,
'fluid_rest': 0.05,
'damping': 0.01,
'subsetp_count': 5,
'iteration_count': 8,
'buoyancy': 1.0})
commands.append(create_container)
if self.use_fluid:
commands.append({"$type": "set_time_step", "time_step": 0.005})
return commands
def get_trial_initialization_commands(self) -> List[dict]:
# clear the flex data
FlexDataset.get_trial_initialization_commands(self)
return Dominoes.get_trial_initialization_commands(self)
def _get_send_data_commands(self) -> List[dict]:
commands = Dominoes._get_send_data_commands(self)
commands.extend(FlexDataset._get_send_data_commands(self))
return commands
def add_rigid_physics_object(self, *args, **kwargs):
"""
Make sure controller knows to treat probe, zone, target, etc. as non-flex objects
"""
o_id = kwargs.get('o_id', None)
if o_id is None:
o_id: int = self.get_unique_id()
kwargs['o_id'] = o_id
commands = Dominoes.add_physics_object(self, *args, **kwargs)
self.non_flex_objects.append(o_id)
print("Add rigid physics object", o_id)
return commands
def add_flex_solid_object(self,
record: ModelRecord,
position: Dict[str, float],
rotation: Dict[str, float],
mesh_expansion: float = 0,
particle_spacing: float = 0.035,
mass: float = 1,
scale: Optional[Dict[str, float]] = {"x": 0.1, "y": 0.5, "z": 0.25},
material: Optional[str] = None,
color: Optional[list] = None,
exclude_color: Optional[list] = None,
o_id: Optional[int] = None,
add_data: Optional[bool] = True,
**kwargs) -> List[dict]:
# so objects don't get stuck in each other -- an unfortunate feature of FLEX
position = {'x': position['x'], 'y': position['y'] + 0.1, 'z': position['z']}
commands = FlexDataset.add_solid_object(
self,
record = record,
position = position,
rotation = rotation,
scale = scale,
mesh_expansion = mesh_expansion,
particle_spacing = particle_spacing,
mass_scale = 1,
o_id = o_id)
# set mass
commands.append({"$type": "set_flex_object_mass",
"mass": mass,
"id": o_id})
# set material and color
commands.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)
commands.append(
{"$type": "set_color",
"color": {"r": color[0], "g": color[1], "b": color[2], "a": 1.},
"id": o_id})
# step physics
if bool(self.step_physics):
print("stepping physics forward", self.step_physics)
commands.append({"$type": "step_physics",
"frames": self.step_physics})
# add data
print("Add FLEX physics object", o_id)
if add_data:
self._add_name_scale_color(record, {'color': color, 'scale': scale, 'id': o_id})
self.masses = np.append(self.masses, mass)
return commands
# def _place_and_push_probe_object(self):
# return []
def _get_push_cmd(self, o_id, position_or_particle=None):
if not self.all_flex_objects:
return Dominoes._get_push_cmd(self, o_id, position_or_particle)
cmd = {"$type": "apply_force_to_flex_object",
"force": self.push_force,
"id": o_id,
"particle": -1}
print("PUSH CMD FLEX")
print(cmd)
return cmd
def drop_cloth(self) -> List[dict]:
self.cloth = self.CLOTH_RECORD
self.cloth_id = self._get_next_object_id()
self.cloth_position = copy.deepcopy({'x':1.0, 'y':1.5,'z':0.0})
self.cloth_color = [0.8,0.5,1.0]
self.cloth_scale = {'x': 1.0, 'y': 1.0, 'z': 1.0}
self.cloth_mass = 0.5
commands = self.add_cloth_object(
record = self.cloth,
position = self.cloth_position,
rotation = {k:0 for k in ['x','y','z']},
scale=self.cloth_scale,
mass_scale = 1,
mesh_tesselation = 1,
tether_stiffness = 1.,
bend_stiffness = 1.,
stretch_stiffness = 1.,
o_id = self.cloth_id)
# set mass
commands.append({"$type": "set_flex_object_mass",
"mass": self.cloth_mass,
"id": self.cloth_id})
# color cloth
commands.append(
{"$type": "set_color",
"color": {"r": self.cloth_color[0], "g": self.cloth_color[1], "b": self.cloth_color[2], "a": 1.},
"id": self.cloth_id})
self._add_name_scale_color(
self.cloth, {'color': self.cloth_color, 'scale': self.cloth_scale, 'id': self.cloth_id})
self.masses = np.append(self.masses, self.cloth_mass)
return commands
def drop_squishy(self) -> List[dict]:
self.squishy = self.SOFT_RECORD
self.squishy_id = self._get_next_object_id()
self.squishy_position = {'x': 0., 'y': 1.0, 'z': 0.}
rotation = {k:0 for k in ['x','y','z']}
self.squishy_color = [0.0,0.8,1.0]
self.squishy_scale = get_random_xyz_transform(self.middle_scale_range)
self.squishy_mass = 2.0
commands = self.add_soft_object(
record = self.squishy,
position = self.squishy_position,
rotation = rotation,
scale=self.squishy_scale,
o_id = self.squishy_id)
# set mass
commands.append({"$type": "set_flex_object_mass",
"mass": self.squishy_mass,
"id": self.squishy_id})
commands.append(
{"$type": "set_color",
"color": {"r": self.squishy_color[0], "g": self.squishy_color[1], "b": self.squishy_color[2], "a": 1.},
"id": self.squishy_id})
self._add_name_scale_color(
self.squishy, {'color': self.squishy_color, 'scale': self.squishy_scale, 'id': self.squishy_id})
self.masses = np.append(self.masses, self.squishy_mass)
return commands
def drop_fluid(self) -> List[dict]:
commands = []
# create a pool for the fluid
self.pool_id = self._get_next_object_id()
print("POOL ID", self.pool_id)
self.non_flex_objects.append(self.pool_id)
commands.append(self.add_transforms_object(record=self.RECEPTACLE_RECORD,
position=TDWUtils.VECTOR3_ZERO,
rotation=TDWUtils.VECTOR3_ZERO,
o_id=self.pool_id,
add_data=True))
commands.append({"$type": "set_kinematic_state",
"id": self.pool_id,
"is_kinematic": True,
"use_gravity": False})
# add the fluid; this will also step physics forward 500 times
self.fluid_id = self._get_next_object_id()
print("FLUID ID", self.fluid_id)
commands.extend(self.add_fluid_object(
position={"x": 0.0, "y": 1.0, "z": 0.0},
rotation=TDWUtils.VECTOR3_ZERO,
o_id=self.fluid_id,
fluid_type=self.ft_selection))
self.fluid_object_ids.append(self.fluid_id)
# restore usual time step
commands.append({"$type": "set_time_step", "time_step": 0.01})
return commands
def _place_ramp_under_probe(self) -> List[dict]:
cmds = Dominoes._place_ramp_under_probe(self)
self.non_flex_objects.append(self.ramp_id)
if self.ramp_base_height >= 0.01:
self.non_flex_objects.append(self.ramp_base_id)
return cmds
def _build_intermediate_structure(self) -> List[dict]:
commands = []
# step physics
# if self.all_flex_objects:
# commands.append({"$type": "step_physics",
# "frames": 50})
commands.extend(self.drop_fluid() if self.use_fluid else [])
commands.extend(self.drop_cloth() if self.use_cloth else [])
commands.extend(self.drop_squishy() if self.use_squishy else [])
return commands
from pathlib import Path
from tdw.controller import Controller
from tdw.tdw_utils import TDWUtils
from tdw.librarian import ModelLibrarian
from tdw.output_data import OutputData, Bounds, Images
"""
1. Add a table and place an object on the table.
2. Add a camera and receive an image.
"""
lib = ModelLibrarian("models_core.json")
# Get the record for the table.
table_record = lib.get_record("small_table_green_marble")
c = Controller()
table_id = 0
# 1. Load the scene.
# 2. Create an empty room (using a wrapper function)
# 3. Add the table.
# 4. Request Bounds data.
resp = c.communicate([{
"$type": "load_scene",
"scene_name": "ProcGenScene"
},
TDWUtils.create_empty_room(12, 12),
c.get_add_object(model_name=table_record.name,
object_id=table_id,
position={
"x": 0,
def __init__(self, port: int = 1071):
lib = ModelLibrarian(str(Path("toys.json").resolve()))
self.records = lib.records
self._target_id: int = 0
super().__init__(port=port)
def run(self):
self.start()
init_setup_commands = [{
"$type": "set_screen_size",
"width": 600,
"height": 480
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
self.communicate(init_setup_commands)
# Create an empty room.
self.communicate({
"$type": "create_empty_environment",
"center": {
"x": 0,
"y": 0,
"z": 0
},
"bounds": {
"x": 8,
"y": 8,
"z": 8
}
})
self.communicate({"$type": "set_gravity", "value": False})
cube_record = ModelLibrarian("models_special.json").get_record(
"prim_cube")
self.communicate({
"$type": "add_object",
"name": "prim_cube",
"url": cube_record.get_url(),
"scale_factor": cube_record.scale_factor,
"position": {
"x": 0,
"y": 0,
"z": 0
},
"rotation": {
"x": 0,
"y": 0,
"z": 0
},
"category": cube_record.wcategory,
"id": 1
})
self.communicate({
"$type": "scale_object",
"scale_factor": {
"x": 30,
"y": 0.0001,
"z": 30
},
"id": 1
})
# Add the avatar.
self.communicate(
TDWUtils.create_avatar(position={
"x": 0,
"y": 0.6,
"z": 0
},
look_at=TDWUtils.array_to_vector3(
[0.5, 0.5, 0]),
avatar_id="avatar"))
self.communicate(self.get_add_hdri_skybox("table_mountain_1_4k"))
self.communicate({"$type": "rotate_hdri_skybox_by", "angle": 90})
lib = MaterialLibrarian(library="materials_med.json")
record = lib.get_record("bricks_chatham_gray_used")
self.communicate({
"$type": "add_material",
"name": "bricks_chatham_gray_used",
"url": record.get_url()
})
self.communicate(
TDWUtils.set_visual_material(c=self,
substructure=cube_record.substructure,
object_id=1,
material="bricks_chatham_gray_used"))
# self.communicate({"$type": "set_field_of_view",
# "field_of_view": 68.0,
# "avatar_id": "avatar"})
bench = self.add_object(model_name="b04_wood_metal_park_bench",
position={
"x": 2,
"y": 0,
"z": 0.5
},
rotation={
"x": 0,
"y": -90,
"z": 0
},
library="models_full.json")
self.add_object(model_name="b04_wood_metal_park_bench",
position={
"x": 5,
"y": 0,
"z": 0.5
},
rotation={
"x": 0,
"y": -90,
"z": 0
},
library="models_full.json")
bench_bounds = self.get_bounds_data(bench)
top = bench_bounds.get_top(0)
self.add_object(model_name="cgaxis_models_65_06_vray",
position={
"x": 1.8,
"y": top[1] - 0.42,
"z": 0.35
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library="models_full.json")
# Enable image capture
self.communicate({
"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]
})
self.communicate({"$type": "send_images", "frequency": "always"})
scene_data = self.communicate({
"$type":
"look_at_position",
"avatar_id":
"avatar",
"position":
TDWUtils.array_to_vector3([0.5, 0.5, 0])
})
images = Images(scene_data[0])
TDWUtils.save_images(
images,
"bench_book",
output_directory=
"/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/compare_COCO_TDW/replicated_images/exterior"
)
def run(self):
self.start()
init_setup_commands = [{"$type": "set_screen_size",
"width": 640,
"height": 480},
{"$type": "set_render_quality",
"render_quality": 5}]
self.communicate(init_setup_commands)
self.communicate({"$type": "create_empty_environment",
"center": {"x": 0, "y": 0, "z": 0},
"bounds": {"x": 15, "y": 15, "z": 15}})
# Disable physics.
self.communicate({"$type": "set_gravity",
"value": False})
# Add the avatar.
self.communicate(TDWUtils.create_avatar(position={"x": -1.5, "y": 1.4, "z": 0.6},
look_at=TDWUtils.array_to_vector3([3, 0.5, 0.5]),
avatar_id="avatar"))
self.communicate(self.get_add_hdri_skybox("table_mountain_1_4k"))
self.communicate({"$type": "set_field_of_view",
"field_of_view": 68.0,
"avatar_id": "avatar"})
cube_record = ModelLibrarian("models_special.json").get_record("prim_cube")
self.communicate({"$type": "add_object",
"name": "prim_cube",
"url": cube_record.get_url(),
"scale_factor": cube_record.scale_factor,
"position": {"x": 0, "y": 0, "z": 0},
"rotation": {"x": 0, "y": 0, "z": 0},
"category": cube_record.wcategory,
"id": 1})
self.communicate({"$type": "scale_object",
"scale_factor": {"x": 30, "y": 0.0001, "z": 30},
"id": 1})
grass_record = MaterialLibrarian("materials_high.json").get_record("grass_countryside")
self.communicate({"$type": "add_material", "name": "grass_countryside", "url": grass_record.get_url()})
self.communicate(TDWUtils.set_visual_material(c=self, substructure=cube_record.substructure, object_id=1,
material="grass_countryside"))
self.add_object(model_name="b03_horse",
position={"x": 3, "y": 0.0001, "z": 0.5},
rotation={"x": 0, "y": 0, "z": 0},
library="models_full.json")
bird = self.add_object(model_name="polysurface63",
position={"x": 1.2, "y": 0.0001, "z": 1},
rotation={"x": 0, "y": 0, "z": 0},
library="models_full.json")
self.communicate({"$type": "scale_object",
"scale_factor": {"x": 2, "y": 2, "z": 2},
"id": bird})
bird = self.add_object(model_name="b03_realistic_pigeon_max",
position={"x": 1.8, "y": 0.0001, "z": 0.45},
rotation={"x": 0, "y": -95, "z": 0},
library="models_full.json")
self.communicate({"$type": "scale_object",
"scale_factor": {"x": 2, "y": 2, "z": 2},
"id": bird})
bird = self.add_object(model_name="rockdove_polysurface77",
position={"x": 1, "y": 0.0001, "z": 0.7},
rotation={"x": 0, "y": -80, "z": 0},
library="models_full.json")
self.communicate({"$type": "scale_object",
"scale_factor": {"x": 2, "y": 2, "z": 2},
"id": bird})
# Enable image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
scene_data = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": TDWUtils.array_to_vector3([3, 0.5, 0])})
images = Images(scene_data[0])
TDWUtils.save_images(images, "bird2",
output_directory="/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/compare_COCO_TDW/replicated_images/exterior")