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 _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):
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 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, 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 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 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 __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 __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 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
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 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")
class RubeGoldbergDemo(Controller):
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 run(self, num_trials: int):
"""
Build a "Rube Goldberg" machine to produce impact sounds.
"""
# Load the photorealistic "archviz_house" environment.
self.load_streamed_scene(scene="archviz_house_2018")
# Organize all initialization commands into a single list.
# Set global values, including the desired screen size and aspect ratio (720P).
# Frame rate is set to 60 fps to facilitate screen / video recording.
init_commands = [{"$type": "set_render_quality",
"render_quality": 5},
{"$type": "set_screen_size",
"width": 1280,
"height": 720},
{"$type": "set_target_framerate",
"framerate": 60},
{"$type": "set_time_step",
"time_step": 0.02}]
# Create the avatar.
init_commands.extend(TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
position={"x": -15.57, "y": 1.886, "z": -4.97}))
# Aim the avatar camera to frame the desired view.
init_commands.extend([{"$type": "rotate_sensor_container_by",
"axis": "yaw",
"angle": 109.13,
"sensor_name": "SensorContainer",
"avatar_id": "a"},
{"$type": "rotate_sensor_container_by",
"axis": "pitch",
"angle": 6.36,
"sensor_name": "SensorContainer",
"avatar_id": "a"}])
# Add the audio sensor.
init_commands.extend([{"$type": "add_audio_sensor",
"avatar_id": "a"}])
# Adjust post-processing settings.
init_commands.extend([{"$type": "set_post_exposure",
"post_exposure": 0.35},
{"$type": "set_screen_space_reflections",
"enabled": True},
{"$type": "set_vignette",
"enabled": False},
{"$type": "set_ambient_occlusion_intensity",
"intensity": 0.175},
{"$type": "set_ambient_occlusion_thickness_modifier",
"thickness": 5.0}])
# Set the shadow strength to maximum.
init_commands.extend([{"$type": "set_shadow_strength",
"strength": 1.0}])
# Send all of the initialization commands.
self.communicate(init_commands)
for i in range(num_trials):
self.do_trial()
def do_trial(self):
# Initialize PyImpact and pass in the "master gain" amplitude value. This value must be betweem 0 and 1.
# The relative amplitudes of all scene objects involved in collisions will be scaled relative to this value.
# Note -- if this value is too high, waveform clipping can occur and the resultant audio will be distorted.
# For this reason, the value used here is considerably smaller than the corresponding value used in the
# impact_sounds.py example controller. Here we have a large number of closely-occuring collisions resulting in
# a rapid series of "clustered" impact sounds, as opposed to a single object falling from a height;
# using a higher value such as the 0.5 used in the example controller will definitely result in unpleasant
# distortion of the audio.
# Note that logging is also enabled.
# Keep track of trial number.
self.current_trial_num += 1
# Create folder for this trial's logging info.
dest_dir = self.root_dest_dir.joinpath(str(self.current_trial_num))
if not dest_dir.exists():
dest_dir.mkdir(parents=True)
dest_dir_str = str(dest_dir.resolve())
p = PyImpact(0.25, logging=True)
self.add_all_objects()
# "Aim" the ball at the monkey and apply the force.
# Note that this force value was arrived at through a number of trial-and-error iterations.
resp = self.communicate([{"$type": "object_look_at_position",
"id": 0,
"position": {"x": -12.95, "y": 1.591, "z": -5.1}},
{"$type": "apply_force_magnitude_to_object",
"id": 0,
"magnitude": 98.0}])
for i in range(400):
collisions, environment_collision, rigidbodies = PyImpact.get_collisions(resp)
# Sort the objects by mass.
masses: Dict[int, float] = {}
for j in range(rigidbodies.get_num()):
masses.update({rigidbodies.get_id(j): rigidbodies.get_mass(j)})
# If there was a collision, create an impact sound.
if len(collisions) > 0 and PyImpact.is_valid_collision(collisions[0]):
collider_id = collisions[0].get_collider_id()
collidee_id = collisions[0].get_collidee_id()
# The "target" object should have less mass than the "other" object.
if masses[collider_id] < masses[collidee_id]:
target_id = collider_id
other_id = collidee_id
else:
target_id = collidee_id
other_id = collider_id
target_name = self.obj_name_dict[target_id]
other_name = self.obj_name_dict[other_id]
impact_sound_command = p.get_impact_sound_command(
collision=collisions[0],
rigidbodies=rigidbodies,
target_id=target_id,
target_mat=self.object_audio_data[target_name].material,
other_id=other_id,
other_mat=self.object_audio_data[other_name].material,
target_amp=self.object_audio_data[target_name].amp,
other_amp=self.object_audio_data[other_name].amp,
resonance=self.object_audio_data[target_name].resonance)
resp = self.communicate(impact_sound_command)
# Continue to run the trial.
else:
resp = self.communicate([])
# Get the logging info for this trial and write it out.
log = p.get_log()
json_dest = dest_dir.joinpath("mode_properties_log.json")
json_dest.write_text(json.dumps(log, indent=2))
for obj_setup in self.object_setups:
self.communicate({"$type": "destroy_object", "id": obj_setup.id})
def add_all_objects(self):
object_commands = []
rigidbodies = []
# Set the mass and scale of the objects, from the data files we parsed earlier.
# Enable output of collision and rigid body data.
# Build dictionary of id,,name so we can retrieve object names during collisions.
# Cache the ids for objects we need to change materials for later
for obj_setup in self.object_setups:
self.obj_name_dict[obj_setup.id] = obj_setup.model_name
rigidbodies.append(obj_setup.id)
if obj_setup.model_name == "prim_sphere":
ball_id = obj_setup.id
elif obj_setup.model_name == "102_pepsi_can_12_fl_oz_vray":
coke_can_id = obj_setup.id
elif obj_setup.model_name == "wood_board":
board_id = obj_setup.id
elif obj_setup.model_name == "bench":
table_id = obj_setup.id
elif obj_setup.model_name == "camera_box":
camera_box_id = obj_setup.id
# Set up to add all objects
object_commands.extend([self.get_add_object(
model_name=obj_setup.model_name,
object_id=obj_setup.id,
position=obj_setup.position,
rotation=obj_setup.rotation,
library=self.object_audio_data[obj_setup.model_name].library),
{"$type": "set_mass",
"id": obj_setup.id,
"mass": self.object_audio_data[obj_setup.model_name].mass},
{"$type": "set_object_collision_detection_mode",
"id": obj_setup.id,
"mode": "continuous_speculative"},
{"$type": "scale_object",
"id": obj_setup.id,
"scale_factor": obj_setup.scale}])
object_commands.extend([{"$type": "send_collisions",
"enter": True,
"stay": False,
"exit": False},
{"$type": "send_rigidbodies",
"frequency": "always",
"ids": rigidbodies}])
# Scale the ball and set a suitable drag value to "tune" how hard it will hit the monkey.
# This combined with the force applied below hits the monkey just hard enough to set the sequnce of
# collisions in motion.
object_commands.extend([{"$type": "scale_object",
"id": ball_id,
"scale_factor": {"x": 0.1, "y": 0.1, "z": 0.1}},
{"$type": "set_object_drag",
"angular_drag": 5.0,
"drag": 1.0,
"id": ball_id}])
# Set physics material parameters to enable rolling motion for the coke can and board after being hit,
# so they collide with the row of "dominos"
object_commands.extend([{"$type": "set_physic_material",
"dynamic_friction": 0.4,
"static_friction": 0.4,
"bounciness": 0.6,
"id": table_id},
{"$type": "set_physic_material",
"dynamic_friction": 0.2,
"static_friction": 0.2,
"bounciness": 0.6,
"id": coke_can_id},
{"$type": "set_physic_material",
"dynamic_friction": 0.2,
"static_friction": 0.2,
"bounciness": 0.7,
"id": board_id}])
# Set the camera box and table to be kinematic, we don't need them to respond to physics.
object_commands.extend([{"$type": "set_kinematic_state",
"id": table_id,
"is_kinematic": True},
{"$type": "set_kinematic_state",
"id": camera_box_id,
"is_kinematic": True}])
# Set the visual material of the ball to metal and the board to a different wood than the bench.
object_commands.extend(TDWUtils.set_visual_material(self, self.ball_record.substructure, ball_id,
"dmd_metallic_fine", quality="high"))
object_commands.extend(TDWUtils.set_visual_material(self, self.board_record.substructure, board_id,
"wood_tropical_hardwood", quality="high"))
# Send all of the object setup commands.
self.communicate(object_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 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,
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
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 __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)
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)
}
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)
class Controller(object):
"""
Base class for all controllers.
Usage:
```python
from tdw.controller import Controller
c = Controller()
c.start()
```
"""
def __init__(self,
port: int = 1071,
check_version: bool = True,
launch_build: bool = True):
"""
Create the network socket and bind the socket to the port.
:param port: The port number.
:param check_version: If true, the controller will check the version of the build and print the result.
:param launch_build: If True, automatically launch the build. If one doesn't exist, download and extract the correct version. Set this to False to use your own build, or (if you are a backend developer) to use Unity Editor.
"""
# Compare the installed version of the tdw Python module to the latest on PyPi.
# If there is a difference, recommend an upgrade.
if check_version:
self._check_pypi_version()
# Launch the build.
if launch_build:
Controller.launch_build()
context = zmq.Context()
self.socket = context.socket(zmq.REP)
self.socket.bind('tcp://*:' + str(port))
self.socket.recv()
self.model_librarian: Optional[ModelLibrarian] = None
self.scene_librarian: Optional[SceneLibrarian] = None
self.material_librarian: Optional[MaterialLibrarian] = None
self.hdri_skybox_librarian: Optional[HDRISkyboxLibrarian] = None
self.humanoid_librarian: Optional[HumanoidLibrarian] = None
self.humanoid_animation_librarian: Optional[
HumanoidAnimationLibrarian] = None
# Compare the version of the tdw module to the build version.
if check_version and launch_build:
self._check_build_version()
def communicate(self, commands: Union[dict, List[dict]]) -> list:
"""
Send commands and receive output data in response.
:param commands: A list of JSON commands.
:return The output data from the build.
"""
if not isinstance(commands, list):
commands = [commands]
self.socket.send_multipart([json.dumps(commands).encode('utf-8')])
return self.socket.recv_multipart()
def start(self, scene="ProcGenScene") -> None:
"""
Init TDW.
:param scene: The scene to load.
"""
self.communicate([{"$type": "load_scene", "scene_name": scene}])
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 get_add_material(self, material_name: str, library: str = "") -> dict:
"""
Returns a valid add_material command.
:param material_name: The name of the material.
:param library: The path to the records file. If left empty, the default library will be selected. See `MaterialLibrarian.get_library_filenames()` and `MaterialLibrarian.get_default_library()`.
:return An add_material command that the controller can then send.
"""
if self.material_librarian is None:
self.material_librarian = MaterialLibrarian(library=library)
record = self.material_librarian.get_record(material_name)
return {
"$type": "add_material",
"name": material_name,
"url": record.get_url()
}
def get_add_scene(self, scene_name: str, library: str = "") -> dict:
"""
Returns a valid add_scene command.
:param scene_name: The name of the scene.
:param library: The path to the records file. If left empty, the default library will be selected. See `SceneLibrarian.get_library_filenames()` and `SceneLibrarian.get_default_library()`.
:return An add_scene command that the controller can then send.
"""
if self.scene_librarian is None:
self.scene_librarian = SceneLibrarian(library=library)
record = self.scene_librarian.get_record(scene_name)
return {
"$type": "add_scene",
"name": scene_name,
"url": record.get_url()
}
def get_add_hdri_skybox(self, skybox_name: str, library: str = "") -> dict:
"""
Returns a valid add_hdri_skybox command.
:param skybox_name: The name of the skybox.
:param library: The path to the records file. If left empty, the default library will be selected. See `HDRISkyboxLibrarian.get_library_filenames()` and `HDRISkyboxLibrarian.get_default_library()`.
:return An add_hdri_skybox command that the controller can then send.
"""
if self.hdri_skybox_librarian is None:
self.hdri_skybox_librarian = HDRISkyboxLibrarian(library=library)
record = self.hdri_skybox_librarian.get_record(skybox_name)
return {
"$type": "add_hdri_skybox",
"name": skybox_name,
"url": record.get_url(),
"exposure": record.exposure,
"initial_skybox_rotation": record.initial_skybox_rotation,
"sun_elevation": record.sun_elevation,
"sun_initial_angle": record.sun_initial_angle,
"sun_intensity": record.sun_intensity
}
def get_add_humanoid(self,
humanoid_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_humanoid command.
:param humanoid_name: The name of the humanoid.
:param position: The position of the humanoid.
:param rotation: The starting rotation of the humanoid, in Euler angles.
:param library: The path to the records file. If left empty, the default library will be selected. See `HumanoidLibrarian.get_library_filenames()` and `HumanoidLibrarian.get_default_library()`.
:param object_id: The ID of the new object.
:return An add_humanoid command that the controller can then send.
"""
if self.humanoid_librarian is None or (
library != "" and self.humanoid_librarian.library != library):
self.humanoid_librarian = HumanoidLibrarian(library=library)
record = self.humanoid_librarian.get_record(humanoid_name)
return {
"$type": "add_humanoid",
"name": humanoid_name,
"url": record.get_url(),
"position": position,
"rotation": rotation,
"id": object_id
}
def get_add_humanoid_animation(
self,
humanoid_animation_name: str,
library="") -> (dict, HumanoidAnimationRecord):
"""
Returns a valid add_humanoid_animation command and the record (which you will need to play an animation).
:param humanoid_animation_name: The name of the animation.
:param library: The path to the records file. If left empty, the default library will be selected. See `HumanoidAnimationLibrarian.get_library_filenames()` and `HumanoidAnimationLibrarian.get_default_library()`.
return An add_humanoid_animation command that the controller can then send.
"""
if self.humanoid_animation_librarian is None:
self.humanoid_animation_librarian = HumanoidAnimationLibrarian(
library=library)
record = self.humanoid_animation_librarian.get_record(
humanoid_animation_name)
return {
"$type": "add_humanoid_animation",
"name": humanoid_animation_name,
"url": record.get_url()
}, record
def load_streamed_scene(self, scene="tdw_room_2018") -> None:
"""
Load a streamed scene. This is equivalent to: `c.communicate(c.get_add_scene(scene))`
:param scene: The name of the streamed scene.
"""
self.communicate(self.get_add_scene(scene))
def add_object(self,
model_name: str,
position={
"x": 0,
"y": 0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library: str = "") -> int:
"""
Add a model to the scene. This is equivalent to: `c.communicate(c.get_add_object())`
: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()`.
:return The ID of the new object.
"""
object_id = Controller.get_unique_id()
self.communicate(
self.get_add_object(model_name, object_id, position, rotation,
library))
return object_id
def get_version(self) -> Tuple[str, str]:
"""
Send a send_version command to the build.
:return The TDW version and the Unity Engine version.
"""
resp = self.communicate({"$type": "send_version"})
for r in resp[:-1]:
if Version.get_data_type_id(r) == "vers":
v = Version(r)
return v.get_tdw_version(), v.get_unity_version()
if len(resp) == 1:
raise Exception(
"Tried receiving version output data but didn't receive anything!"
)
raise Exception(f"Expected output data with ID vers but got: " +
Version.get_data_type_id(resp[0]))
@staticmethod
def get_unique_id() -> int:
"""
Generate a unique integer. Useful when creating objects.
:return The new unique ID.
"""
return int.from_bytes(os.urandom(3), byteorder='big')
@staticmethod
def get_frame(frame: bytes) -> int:
"""
Converts the frame byte array to an integer.
:param frame: The frame as bytes.
:return The frame as an integer.
"""
return int.from_bytes(frame, byteorder='big')
@staticmethod
def launch_build() -> None:
"""
Launch the build. If a build doesn't exist at the expected location, download one to that location.
"""
# Download the build.
if not Build.BUILD_PATH.exists():
print(
f"Couldn't find build at {Build.BUILD_PATH}\nDownloading now..."
)
success = Build.download()
if not success:
print("You need to launch your own build.")
else:
success = True
# Launch the build.
if success:
Popen(str(Build.BUILD_PATH.resolve()))
def _check_build_version(self,
version: str = __version__,
build_version: str = None) -> None:
"""
Check the version of the build. If there is no build, download it.
If the build is of the wrong version, recommend an upgrade.
:param version: The version of TDW. You can set this to an arbitrary version for testing purposes.
:param build_version: If not None, this overrides the expected build version. Only override for debugging.
"""
v = PyPi.strip_post_release(version)
tdw_version, unity_version = self.get_version()
# Override the build version for testing.
if build_version is not None:
tdw_version = build_version
pypi_version = PyPi.get_latest_minor_release(tdw_version)
print(
f"Build version {tdw_version}\nUnity Engine {unity_version}\nPython tdw module version {version}"
)
if v < tdw_version:
print(
"WARNING! Your TDW build is newer than your tdw Python module. They might not be compatible."
)
print(
f"To download the correct build:\n\nfrom tdw.release.build import Build\nBuild.download(version={v})"
)
print(
f"\nTo upgrade your Python module (usually recommended):\n\npip3 install tdw=={pypi_version}"
)
elif v > tdw_version:
print(
"WARNING! Your TDW build is older than your tdw Python module. Downloading the correct build..."
)
Build.download(v)
@staticmethod
def _check_pypi_version(v_installed_override: str = None,
v_pypi_override: str = None) -> None:
"""
Compare the version of the tdw Python module to the latest on PyPi.
If there is a mismatch, offer an upgrade recommendation.
:param v_installed_override: Override for the installed version. Change this to debug.
:param v_pypi_override: Override for the PyPi version. Change this to debug.
"""
# Get the version of the installed tdw module.
installed_tdw_version = PyPi.get_installed_tdw_version()
# Get the latest version of the tdw module on PyPi.
pypi_version = PyPi.get_pypi_version()
# Apply overrides
if v_installed_override is not None:
installed_tdw_version = v_installed_override
if v_pypi_override is not None:
pypi_version = v_pypi_override
# If there is a mismatch, recommend an upgrade.
if installed_tdw_version != pypi_version:
# Strip the installed version of the post-release suffix (e.g. 1.6.3.4 to 1.6.3).
stripped_installed_version = PyPi.strip_post_release(
installed_tdw_version)
# This message is here only for debugging.
if stripped_installed_version != __version__:
print(
f"Your installed version: {stripped_installed_version} "
f"doesn't match tdw.version.__version__: {__version__} "
f"(this may be because you're using code from the tdw repo that is ahead of PyPi)."
)
# Strip the latest PyPi version of the post-release suffix.
stripped_pypi_version = PyPi.strip_post_release(pypi_version)
print(
f"You are using TDW {installed_tdw_version} but version {pypi_version} is available."
)
# If user is behind by a post release, recommend an upgrade to the latest.
# (Example: installed version is 1.6.3.4 and PyPi version is 1.6.3.5)
if stripped_installed_version == stripped_pypi_version:
print(
f"Upgrade to the latest version of TDW:\npip3 install tdw -U"
)
# Using a version behind the latest (e.g. latest is 1.6.3 and installed is 1.6.2)
# If the user is behind by a major or minor release, recommend either upgrading to a minor release
# or to a major release.
# (Example: installed version is 1.6.3.4 and PyPi version is 1.7.0.0)
else:
installed_major = PyPi.get_major_release(
stripped_installed_version)
pypi_minor = PyPi.get_latest_minor_release(
stripped_installed_version)
# Minor release mis-match.
if PyPi.strip_post_release(
pypi_minor) != stripped_installed_version:
print(
f"To upgrade to the last version of 1.{installed_major}:\n"
f"pip3 install tdw=={pypi_minor}")
pypi_major = PyPi.get_major_release(stripped_pypi_version)
# Major release mis-match.
if installed_major != pypi_major:
# Offer to upgrade to a major release.
print(
f"Consider upgrading to the latest version of TDW ({stripped_pypi_version}):"
f"\npip3 install tdw -U")
else:
print("Your installed tdw Python module is up to date with PyPi.")