def get_bounds_dict(bounds: Bounds, index: int) -> Dict[str, np.array]:
"""
:param bounds: Bounds output data.
:param index: The index in `bounds` of the target object.
:return: A dictionary of the bounds. Key = the name of the position. Value = the position as a numpy array.
"""
return {
"top": np.array(bounds.get_top(index)),
"bottom": np.array(bounds.get_bottom(index)),
"left": np.array(bounds.get_left(index)),
"right": np.array(bounds.get_right(index)),
"front": np.array(bounds.get_front(index)),
"back": np.array(bounds.get_back(index))
}
def get_bounds_extents(bounds: Bounds, index: int) -> np.array:
"""
:param bounds: Bounds output data.
:param index: The index in `bounds` of the target object.
:return: The width (left to right), length (front to back), and height (top to bottom) of the bounds as a numpy array.
"""
return np.array([
np.linalg.norm(
np.array(bounds.get_left(index)) -
np.array(bounds.get_right(index))),
np.linalg.norm(
np.array(bounds.get_front(index)) -
np.array(bounds.get_back(index))),
np.linalg.norm(
np.array(bounds.get_top(index)) -
np.array(bounds.get_bottom(index)))
])
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)
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) -> None:
self.start()
# Get IDs for the object, the table, and the avatars.
o_id = self.get_unique_id()
table_id = self.get_unique_id()
table_position = {"x": 0, "y": 0, "z": 6.8}
a = "a"
# 1. Create the room.
commands = [TDWUtils.create_empty_room(20, 20)]
# 2. Add the avatar.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_StickyMitten_Adult",
avatar_id="a"))
# 3. Add the objects. Scale the objects. Set a high mass for the table.
# 4. Disable the avatar's cameras.
# 5. Set the stickiness of the avatar's left mitten. Set a high drag. Rotate the head.
# 6. Request AvatarChildrenNames data and Bounds data (this frame only).
# 7. Strengthen the left shoulder a bit.
commands.extend([
self.get_add_object("jug05",
position={
"x": -0.417,
"y": 0.197,
"z": 0.139
},
rotation={
"x": 90,
"y": 0,
"z": 0
},
object_id=o_id), {
"$type": "scale_object",
"scale_factor": {
"x": 2,
"y": 2,
"z": 2
},
"id": o_id
},
self.get_add_object("small_table_green_marble",
position={
"x": 0,
"y": 0,
"z": 6.8
},
rotation={
"x": 0,
"y": 0,
"z": 0
},
object_id=table_id), {
"$type": "scale_object",
"scale_factor": {
"x": 2,
"y": 0.5,
"z": 2
},
"id": table_id
},
{
"$type": "set_object_collision_detection_mode",
"id": table_id,
"mode": "continuous_dynamic"
}, {
"$type": "set_object_collision_detection_mode",
"id": o_id,
"mode": "continuous_dynamic"
}, {
"$type": "set_mass",
"mass": 100,
"id": table_id
}, {
"$type": "toggle_image_sensor",
"sensor_name": "SensorContainer",
"avatar_id": a
}, {
"$type": "toggle_image_sensor",
"sensor_name": "FollowCamera",
"avatar_id": a
}, {
"$type": "set_stickiness",
"sub_mitten": "back",
"sticky": True,
"is_left": True,
"avatar_id": a
}, {
"$type": "set_avatar_drag",
"drag": 1000,
"angular_drag": 1000,
"avatar_id": a
}, {
"$type": "rotate_head_by",
"axis": "pitch",
"angle": 5
}, {
"$type": "send_avatar_children_names",
"ids": [a],
"frequency": "once"
}, {
"$type": "send_bounds",
"ids": [table_id],
"frequency": "once"
}, {
"$type": "adjust_joint_force_by",
"delta": 2,
"joint": "shoulder_left",
"axis": "pitch"
}
])
# 7. Add a 3rd-person camera.
commands.extend(
TDWUtils.create_avatar(avatar_type="A_Img_Caps_Kinematic",
avatar_id="c",
position={
"x": -3.9,
"y": 2.3,
"z": 4.3
}))
# 8. Request StickyMittenAvatar and Images data per-frame.
commands.extend([{
"$type": "set_pass_masks",
"pass_masks": ["_img"],
"avatar_id": self.cam_id
}, {
"$type": "send_avatars",
"ids": [a],
"frequency": "always"
}, {
"$type": "send_images",
"ids": [self.cam_id],
"frequency": "always"
}])
resp = self.communicate(commands)
# Get the object ID of the left mitten, the size of the table, and the avatar's position.
mitten_left_id = None
table_size: Tuple[float, float, float] = (0, 0, 0)
for r in resp[:-1]:
r_id = OutputData.get_data_type_id(r)
# Get the mitten ID.
if r_id == "avcn":
avcn = AvatarChildrenNames(r)
for i in range(avcn.get_num_children()):
if avcn.get_child_name(i) == "mitten_left":
mitten_left_id = avcn.get_child_id(i)
# Get the table bounds.
elif r_id == "boun":
boun = Bounds(r)
for i in range(boun.get_num()):
if boun.get_id(i) == table_id:
table_size = (boun.get_right(i)[0] -
boun.get_left(i)[0], boun.get_top(i)[1] -
boun.get_bottom(i)[1],
boun.get_front(i)[2] -
boun.get_back(i)[2])
# Get the avatar's position.
elif r_id == "avsm":
self.avatar_position = AvatarStickyMitten(r).get_position()
# Pick up the object.
# Lift up the object.
self._bend_arm_joints([{
"$type": "pick_up_proximity",
"distance": 20,
"grip": 10000,
"is_left": True,
"avatar_id": a
}, {
"$type": "rotate_head_by",
"axis": "pitch",
"angle": 20,
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 25,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": -25,
"joint": "shoulder_left",
"axis": "yaw",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 60,
"joint": "shoulder_left",
"axis": "roll",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": 100,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Allow the avatar to move again.
self._do_frame({
"$type": "set_avatar_drag",
"drag": 0.125,
"angular_drag": 1000,
"avatar_id": a
})
# Move to the table.
move_to_table = True
# The position of the side of the table the avatar is aiming for.
table_side_position = {
"x": table_position["x"],
"y": 0,
"z": table_position["z"] - table_size[2]
}
while move_to_table:
# Stop moving if we are close enough.
if TDWUtils.get_distance(
table_side_position,
TDWUtils.array_to_vector3(self.avatar_position)) < 0.7:
move_to_table = False
# Keep moving forward.
else:
self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": 20
})
# Stop.
# Allow the avatar to move again.
self._do_frame({
"$type": "set_avatar_drag",
"drag": 1000,
"angular_drag": 1000,
"avatar_id": a
})
mitten_over_table = False
target_z = table_side_position["z"] + 0.17
# Keep lifting until the mitten is over the table.
while not mitten_over_table:
avsm = self._bend_arm_joints([{
"$type": "bend_arm_joint_by",
"angle": 15,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_by",
"angle": -10,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten is over the table.
mitten_over_table = avsm.get_body_part_position(
i)[2] >= target_z
# Drop the arm gently.
mitten_near_table = False
while not mitten_near_table:
avsm = self._bend_arm_joints([{
"$type": "bend_arm_joint_by",
"angle": -25,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}])
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten very close to the table surface.
mitten_near_table = avsm.get_body_part_position(
i)[1] - table_size[1] <= 0.1
# Drop the object.
# Allow the avatar to move again.
self._do_frame([{
"$type": "put_down",
"is_left": True,
"avatar_id": a
}, {
"$type": "set_avatar_drag",
"drag": 0.125,
"angular_drag": 1000,
"avatar_id": a
}])
# Back away from the table.
mitten_away_from_table = False
while not mitten_away_from_table:
avsm = self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": -50
})
# Get the mitten.
for i in range(avsm.get_num_body_parts()):
if avsm.get_body_part_id(i) == mitten_left_id:
# Check if the mitten has cleared the table.
mitten_away_from_table = avsm.get_body_part_position(
i)[2] < target_z
# Drop the arm.
self._do_frame([{
"$type": "rotate_head_by",
"axis": "pitch",
"angle": -25
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "pitch",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "yaw",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "shoulder_left",
"axis": "roll",
"avatar_id": a
}, {
"$type": "bend_arm_joint_to",
"angle": 0,
"joint": "elbow_left",
"axis": "pitch",
"avatar_id": a
}])
# Back away from the table.
away_from_table = False
while not away_from_table:
self._do_frame({
"$type": "move_avatar_forward_by",
"avatar_id": a,
"magnitude": -20
})
away_from_table = TDWUtils.get_distance(
TDWUtils.array_to_vector3(self.avatar_position),
table_position) > 1.8
# Let the joints drop.
self._bend_arm_joints([])
self.communicate({"$type": "terminate"})
def __init__(self,
object_id: int,
rigidbodies: Rigidbodies,
segmentation_colors: SegmentationColors,
bounds: Bounds,
audio: ObjectInfo,
target_object: bool = False):
"""
:param object_id: The unique ID of the object.
:param rigidbodies: Rigidbodies output data.
:param bounds: Bounds output data.
:param segmentation_colors: Segmentation colors output data.
"""
self.object_id = object_id
self.model_name = audio.name
self.container = self.model_name in StaticObjectInfo.CONTAINERS
self.kinematic = self.model_name in StaticObjectInfo._KINEMATIC
self.target_object = target_object
self.category = ""
# This is a sub-object of a composite object.
if audio.library == "":
# Get the record of the composite object.
for k in StaticObjectInfo._COMPOSITE_OBJECTS:
for v in StaticObjectInfo._COMPOSITE_OBJECTS[k]:
if v == audio.name:
record = TransformInitData.LIBRARIES[
"models_core.json"].get_record(k)
# Get the semantic category.
self.category = record.wcategory
break
else:
# Get the model record from the audio data.
record = TransformInitData.LIBRARIES[audio.library].get_record(
audio.name)
# Get the semantic category.
self.category = record.wcategory
# Get the segmentation color.
self.segmentation_color: Optional[np.array] = None
for i in range(segmentation_colors.get_num()):
if segmentation_colors.get_object_id(i) == self.object_id:
self.segmentation_color = np.array(
segmentation_colors.get_object_color(i))
break
assert self.segmentation_color is not None, f"Segmentation color not found: {self.object_id}"
# Get the size of the object.
self.size = np.array([0, 0, 0])
for i in range(bounds.get_num()):
if bounds.get_id(i) == self.object_id:
self.size = np.array([
float(
np.abs(bounds.get_right(i)[0] -
bounds.get_left(i)[0])),
float(
np.abs(bounds.get_top(i)[1] -
bounds.get_bottom(i)[1])),
float(
np.abs(bounds.get_front(i)[2] - bounds.get_back(i)[2]))
])
break
assert np.linalg.norm(
self.size) > 0, f"Bounds data not found for: {self.object_id}"
# Get the mass.
self.mass: float = -1
for i in range(rigidbodies.get_num()):
if rigidbodies.get_id(i) == self.object_id:
self.mass = rigidbodies.get_mass(i)
break
assert self.mass >= 0, f"Mass not found: {self.object_id}"