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 __init__(self, root_dir: Path):
animation_lib = HumanoidAnimationLibrarian()
scene_lib = SceneLibrarian()
animation_scene_matrix = json.loads(
Path("animation_scene_matrix.json").read_text())
self.animations_and_scenes: List[_AnimationAndScene] = []
for a in animation_scene_matrix:
for s in animation_scene_matrix[a]:
combo = _AnimationAndScene(
animation=animation_lib.get_record(a),
scene=scene_lib.get_record(s),
position=animation_scene_matrix[a][s]["position"],
rotation=animation_scene_matrix[a][s]["rotation"])
self.animations_and_scenes.append(combo)
super().__init__()
self.root_dir = root_dir
if not self.root_dir.exists():
self.root_dir.mkdir(parents=True)
# Set global values.
self.communicate([{
"$type": "set_render_quality",
"render_quality": 5
}, {
"$type": "set_img_pass_encoding",
"value": False
}, {
"$type": "set_vignette",
"enabled": False
}, {
"$type": "set_shadow_strength",
"strength": 1.0
}, {
"$type": "set_screen_size",
"width": 1280,
"height": 720
}])
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.")
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):
# Setting up image directory
output_directory = "multiple_objects_pics"
parent_dir = '/Python/Leonard/multiple_objects'
path = os.path.join(parent_dir, output_directory)
print(path)
if os.path.exists(path):
shutil.rmtree(path)
time.sleep(0.5)
os.mkdir(path)
# Initialize scene
self.start()
# Load in scene
scene_lib = SceneLibrarian()
scene_record = scene_lib.get_record("tdw_room_2018")
self.communicate({"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()})
# Resize screen; note render_quality 5 is best
self.communicate([{"$type": "set_screen_size",
"width": 1280,
"height": 962},
{"$type": "set_render_quality",
"render_quality": 2}])
# Toggle physics
physics_on = True
self.communicate({"$type": "simulate_physics",
"value": physics_on})
# Add in coffee table object
table_id = self.add_object("live_edge_coffee_table",
position={"x": 1.5, "y": 0, "z": 1.5},
rotation={"x": 0, "y": 0, "z": 0})
# Add two bowls on top of the table
bowl_1 = self.add_object(model_name="b04_bowl_smooth",
position={"x": 1.25, "y": 0.331, "z": 1.5},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_full.json")
self.add_object(model_name="b04_bowl_smooth",
position={"x": 1.75, "y": 0.331, "z": 1.5},
rotation=TDWUtils.VECTOR3_ZERO,
library="models_full.json")
# Push one bowl towards the other
self.communicate({"$type": "apply_force_to_object",
"force": {"x": 5, "y": 0, "z": 0},
"id": bowl_1})
# Load in avatar
self.communicate(TDWUtils.create_avatar(avatar_id="avatar",
position={"x": 0, "y": 1.5, "z": 0},
look_at={"x": 1.5, "y": 0, "z": 1.5}))
# Necessary commands to avoid error in calling save_images()
# Second command enables image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
# Capture images
for i in range(20):
resp = self.communicate({"$type": "look_at",
"avatar_id": "avatar",
"object_id": table_id,
"use_centroid": True})
images = Images(resp[0])
TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=path)
def run(self):
# Setting up image directory
output_directory = "rotating_guitar_images"
parent_dir = '/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/rotating_guitar'
path = os.path.join(parent_dir, output_directory)
print(path)
if os.path.exists(path):
shutil.rmtree(path)
time.sleep(0.5)
os.mkdir(path)
# Initialize scene
self.start()
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()
scene_record = scene_lib.get_record("tdw_room_2018")
self.communicate({
"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()
})
self.communicate({"$type": "set_gravity", "value": True})
self.communicate({"$type": "simulate_physics", "value": True})
self.communicate(
TDWUtils.create_avatar(avatar_id="avatar",
position={
"x": 3,
"y": 1.5,
"z": 3
},
look_at={
"x": 0,
"y": 0.8,
"z": 0
}))
guitar_id = self.add_object(model_name="b01_shovel",
position={
"x": 0,
"y": 1,
"z": 0
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
library="models_full.json")
self.communicate({
"$type": "rotate_object_by",
"angle": -15.0,
"id": guitar_id,
"axis": "roll",
"is_world": True
})
self.communicate({
"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]
})
self.communicate({"$type": "send_images", "frequency": "always"})
for i in range(50):
resp = self.communicate({
"$type": "look_at_position",
"avatar_id": "avatar",
"position": {
"x": 0,
"y": 0.8,
"z": 0
}
})
images = Images(resp[0])
TDWUtils.save_images(images,
TDWUtils.zero_padding(i),
output_directory=path)
static_friction=0.2,
bounciness=0.9,
mass=5),
PhysicsObject(model_name="duffle_bag",
initial_position={
"x": -3,
"y": 2,
"z": 0.5
},
possibility=Possibility(),
dynamic_friction=0.5,
static_friction=0.6,
bounciness=0.9,
mass=8)
]
lib_scenes = SceneLibrarian()
record_scene = lib_scenes.get_record("building_site")
sce_stp = [{
"$type": "add_scene",
"name": record_scene.name,
"url": record_scene.get_url()
}]
stp = Trial(occ_stp, mov_stp, sce_stp, "stp")
# Spatio-Temporal Continuity - Impossible
occ_sti = occ_stp
mov_sti = [
PhysicsObject(model_name="chair_eames_plastic_armchair",
initial_position={
"x": 3,
"y": 2,
def run(self):
# Setting up image directory
output_directory = "floating_bowl_images"
parent_dir = '/Users/leonard/Desktop/TDWBase-1.5.0/Python/Leonard/floating_bowl'
path = os.path.join(parent_dir, output_directory)
print(path)
if os.path.exists(path):
shutil.rmtree(path)
time.sleep(0.5)
os.mkdir(path)
# Initialize scene
self.start()
# Load in scene
scene_lib = SceneLibrarian()
scene_record = scene_lib.get_record("tdw_room_2018")
self.communicate({"$type": "add_scene",
"name": scene_record.name,
"url": scene_record.get_url()})
# Resize screen; note render_quality 5 is best
self.communicate([{"$type": "set_screen_size",
"width": 1280,
"height": 962},
{"$type": "set_render_quality",
"render_quality": 2}])
# Toggle physics
physics_on = True
self.communicate({"$type": "simulate_physics",
"value": physics_on})
# Load in avatar
self.communicate(TDWUtils.create_avatar(avatar_id="avatar",
position={"x": 3, "y": 1.5, "z": 3},
look_at={"x": 0, "y": 0.8, "z": 0}))
# Necessary commands to avoid error in calling save_images()
# Second command enables image capture
self.communicate({"$type": "set_pass_masks",
"avatar_id": "avatar",
"pass_masks": ["_img", "_id"]})
self.communicate({"$type": "send_images",
"frequency": "always"})
# Capture images
for i in range(20):
bowl_id = self.add_object(model_name="b04_bowl_smooth",
position={"x": 0, "y": i / 10, "z": 0},
rotation={"x": i * 30 / 10, "y": i / 10, "z": 0},
library="models_full.json")
resp = self.communicate({"$type": "look_at_position",
"avatar_id": "avatar",
"position": {"x": 0, "y": 0.8, "z": 0}})
images = Images(resp[0])
TDWUtils.save_images(images, TDWUtils.zero_padding(i), output_directory=path)
self.communicate({"$type": "destroy_object",
"id": bowl_id})
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')