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)
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,
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.")
class Submerge(FlexDataset):
"""
Create a fluid "container" with the NVIDIA Flex physics engine.
Run several trials, dropping ball objects of increasing mass into the fluid.
"""
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 get_scene_initialization_commands(self) -> List[dict]:
if system() != "Windows":
raise Exception(
"Flex fluids are only supported in Windows (see Documentation/misc_frontend/flex.md)"
)
commands = [
self.get_add_scene(scene_name="tdw_room_2018"), {
"$type": "set_aperture",
"aperture": 4.8
}, {
"$type": "set_focus_distance",
"focus_distance": 2.25
}, {
"$type": "set_post_exposure",
"post_exposure": 0.4
}, {
"$type": "set_ambient_occlusion_intensity",
"intensity": 0.175
}, {
"$type": "set_ambient_occlusion_thickness_modifier",
"thickness": 3.5
}
]
return commands
def get_trial_initialization_commands(self) -> List[dict]:
super().get_trial_initialization_commands()
trial_commands = []
if self.pool_id is not None:
trial_commands.append({"$type": "destroy_flex_container", "id": 0})
# Load a pool container for the fluid.
self.pool_id = Controller.get_unique_id()
self.non_flex_objects.append(self.pool_id)
trial_commands.append(
self.add_transforms_object(record=self.receptacle_record,
position={
"x": 0,
"y": 0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
o_id=self.pool_id))
trial_commands.extend([{
"$type": "scale_object",
"id": self.pool_id,
"scale_factor": {
"x": 1.5,
"y": 2.5,
"z": 1.5
}
}, {
"$type": "set_kinematic_state",
"id": self.pool_id,
"is_kinematic": True,
"use_gravity": False
}])
# Randomly select a fluid type
fluid_type_selection = choice(self.ft.fluid_type_names)
# Create the container, set up for fluids.
# Slow down physics so the water can settle without splashing out of the container.
trial_commands.extend([{
"$type":
"create_flex_container",
"collision_distance":
0.04,
"static_friction":
0.1,
"dynamic_friction":
0.1,
"particle_friction":
0.1,
"viscocity":
self.ft.fluid_types[fluid_type_selection].viscosity,
"adhesion":
self.ft.fluid_types[fluid_type_selection].adhesion,
"cohesion":
self.ft.fluid_types[fluid_type_selection].cohesion,
"radius":
0.1,
"fluid_rest":
0.05,
"damping":
0.01,
"substep_count":
5,
"iteration_count":
8,
"buoyancy":
1.0
}, {
"$type": "set_time_step",
"time_step": 0.005
}])
# Recreate fluid.
# Add the fluid actor, using the FluidPrimitive. Allow 500 frames for the fluid to settle before continuing.
fluid_id = Controller.get_unique_id()
trial_commands.extend(
self.add_fluid_object(position={
"x": 0,
"y": 1.0,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
o_id=fluid_id,
fluid_type=fluid_type_selection))
trial_commands.append({"$type": "set_time_step", "time_step": 0.03})
# Select an object at random.
model = choice(self.model_list)
model_record = self.full_lib.get_record(model)
info = PHYSICS_INFO[model]
# Randomly select an object, and randomly orient it.
# Set the solid actor and assign the container.
o_id = Controller.get_unique_id()
trial_commands.extend(
self.add_solid_object(record=model_record,
position={
"x": 0,
"y": 2,
"z": 0
},
rotation={
"x": uniform(-45.0, 45.0),
"y": uniform(-45.0, 45.0),
"z": uniform(-45.0, 45.0)
},
o_id=o_id,
scale={
"x": 0.5,
"y": 0.5,
"z": 0.5
},
mass_scale=info.mass,
particle_spacing=0.05))
# Reset physics time-step to a more normal value.
# Position and aim avatar.
trial_commands.extend([{
"$type": "set_kinematic_state",
"id": o_id
}, {
"$type": "teleport_avatar_to",
"position": {
"x": -2.675,
"y": 1.375,
"z": 0
}
}, {
"$type": "look_at",
"object_id": self.pool_id,
"use_centroid": True
}])
return trial_commands
def get_per_frame_commands(self, frame: int,
resp: List[bytes]) -> List[dict]:
return [{"$type": "focus_on_object", "object_id": self.pool_id}]
def get_field_of_view(self) -> float:
return 35
def is_done(self, resp: List[bytes], frame: int) -> bool:
return frame > 200
src_paths = a.prefab_to_asset_bundle(
Path.home().joinpath("asset_bundle_creator/Assets/Resources/prefab"),
model_name=args.name)
src_asset_bundles = dict()
for q in src_paths:
src_asset_bundles[q.parts[-2]] = q
# Parse the URLs.
urls = a.get_local_urls(src_paths)
# Create the metadata record.
record_path = a.create_record(args.name, 2886585, "container", 1, urls)
record = ModelRecord(json.loads(record_path.read_text(encoding="utf-8")))
# Add the record.
r = lib.get_record(record.name)
lib.add_or_update_record(record=record,
overwrite=False if r is None else True,
write=True)
# Make the URLs relative paths.
temp = dict()
for p in record.urls:
dest_dir = f"../asset_bundles/{p}"
dd = root_dest.joinpath(f"asset_bundles/{p}")
if not dd.exists():
dd.mkdir(parents=True)
temp[p] = f"../asset_bundles/{p}/{record.name}"
record.urls = temp
lib.add_or_update_record(record=record, overwrite=True, write=True)
# Copy the asset bundles.
for p in UNITY_TO_SYSTEM:
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):
"""
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.
"""
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:
tdw_version, unity_version = self.get_version()
print(
f"Build version {tdw_version}\nUnity Engine {unity_version}\nPython tdw module version {__version__}"
)
if __version__ != tdw_version:
print(
"WARNING! Your Python code is not the same version as the build. They might not be compatible."
)
print(
"Either use the latest prerelease build, or use the last stable release via:\n"
)
print("git checkout v" + last_stable_release)
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')