def get_image(self, record: ModelRecord):
o_id = Controller.get_unique_id()
self.communicate({"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"rotation": record.canonical_rotation,
"id": o_id})
s = TDWUtils.get_unit_scale(record) * 2
# Scale the model and get an image.
# Look at the model's centroid.
resp = self.communicate([{"$type": "scale_object",
"id": o_id,
"scale_factor": {"x": s, "y": s, "z": s}},
{"$type": "look_at",
"avatar_id": "a",
"object_id": o_id,
"use_centroid": True}])
# Destroy the model and unload the asset bundle.
self.communicate([{"$type": "destroy_object",
"id": o_id},
{"$type": "unload_asset_bundles"}])
return Images(resp[0]), resp[-1]
def add_transforms_object(self,
record: ModelRecord,
position: Dict[str, float],
rotation: Dict[str, float],
o_id: Optional[int] = None) -> dict:
"""
This is a wrapper for `Controller.get_add_object()` and the `add_object` command.
This caches the ID of the object so that it can be easily cleaned up later.
:param record: The model record.
:param position: The initial position of the object.
:param rotation: The initial rotation of the object, in Euler angles.
:param o_id: The unique ID of the object. If None, a random ID is generated.
:return: An `add_object` command.
"""
if o_id is None:
o_id: int = Controller.get_unique_id()
# Log the static data.
self.object_ids = np.append(self.object_ids, o_id)
return {
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": position,
"rotation": rotation,
"category": record.wcategory,
"id": o_id
}
def __init__(self,
record_path: str,
asset_bundle_path: str,
build_path: str,
port=1071):
"""
:param record_path: The path to the temporary record file.
:param asset_bundle_path: The path to the local asset bundle.
:param build_path: The path to the build executable.
:param port: The port.
"""
self.record_path: str = record_path
self.record: ModelRecord = ModelRecord(
json.loads(Path(record_path).read_text(encoding="utf-8")))
self.build_path: str = build_path
self.asset_bundle_path: str = asset_bundle_path
if not self.asset_bundle_path.startswith("file:///"):
self.asset_bundle_path = "file:///" + self.asset_bundle_path
try:
# Create the build.
Popen(build_path)
except:
print("No build found at: " + build_path)
print("You need to launch a build manually.")
super().__init__(port)
def create_library(self) -> ModelLibrarian:
lib = self._get_librarian("ShapeNetSEM")
first_time_only = True
metadata_path = self.src.joinpath("metadata.csv")
with open(str(metadata_path), newline='') as csvfile:
reader = csv.reader(csvfile)
for row in reader:
if first_time_only:
first_time_only = False
continue
if row[1] == "" or row[3] == "":
continue
record = ModelRecord()
record.name = row[0][4:]
record.wcategory = row[3].split(",")[0]
record.wnid = f"n{int(row[2][1:]):08d}"
for platform in record.urls:
record.urls[platform] = self._get_url(
record.wnid, record.name, platform)
lib.add_or_update_record(record, overwrite=False, write=False)
# Write to disk.
lib.write(pretty=False)
# Move the textures.
any_textures = False
for f in self.src.joinpath("textures").rglob("*.jpg"):
f.replace(self.src.joinpath(f"models/{f.name}"))
any_textures = True
if any_textures:
print("Moved all .jpg files in textures/ to models/")
return lib
def __init__(self, record_path: str, asset_bundle_path: str, port=1071):
"""
:param record_path: The path to the temporary record file.
:param asset_bundle_path: The path to the local asset bundle.
:param port: The port.
"""
self.record_path: str = record_path
self.record: ModelRecord = ModelRecord(
json.loads(Path(record_path).read_text(encoding="utf-8")))
self.asset_bundle_path: str = asset_bundle_path
if not self.asset_bundle_path.startswith("file:///"):
self.asset_bundle_path = "file:///" + self.asset_bundle_path
super().__init__(port)
def create_library(self) -> ModelLibrarian:
# Load the taxonomy file.
metadata = json.loads(
self._get_metadata_path().read_text(encoding="utf-8"))
# Create a new library.
lib = self._get_librarian("ShapeNetCoreVal")
# Process each .obj file.
for f in self.src.rglob("*.obj"):
record = ModelRecord()
record.name = f.parts[-3]
record.wnid = metadata[record.name]['wnid']
record.wcategory = metadata[record.name]['wcategory']
for platform in record.urls:
record.urls[platform] = self._get_url(record.wnid, record.name,
platform)
lib.add_or_update_record(record, overwrite=False, write=False)
# Write to disk. Don't pretty-print (saves about 60 MB).
lib.write(pretty=False)
return lib
def create_library(self) -> ModelLibrarian:
# Load the taxonomy file.
taxonomy_raw = json.loads(self._get_taxonomy_path().read_text(encoding="utf-8"))
taxonomy = dict()
for synset in taxonomy_raw:
taxonomy.update({synset["synsetId"]: synset["name"].split(",")[0]})
# Create a new library.
lib = self._get_librarian("ShapeNetCore")
# Process each .obj file.
for f in self.src.rglob("*.obj"):
wnid = f.parts[-4]
record = ModelRecord()
record.name = f.parts[-3]
record.wnid = "n" + wnid
record.wcategory = taxonomy[wnid]
for platform in record.urls:
record.urls[platform] = self._get_url(record.wnid, record.name, platform)
lib.add_or_update_record(record, overwrite=False, write=False)
# Write to disk. Don't pretty-print (saves about 60 MB).
lib.write(pretty=False)
return lib
def create_record(self,
model_name: str,
wnid: int,
wcategory: str,
scale: float,
urls: Dict[str, str],
record: Optional[ModelRecord] = None,
write_physics: bool = False) -> Path:
"""
Create a local .json metadata record of the model.
:param model_name: The name of the model.
:param wnid: The WordNet ID.
:param wcategory: The WordNet category.
:param scale: The default scale of the object.
:param urls: The finalized URLs (or local filepaths) of the assset bundles.
:param record: A pre-written metadata record. If not None, it will override the other parameters.
:param write_physics: If true, launch the build to write the physics quality. (This is optional).
:return The path to the file with the metadata record.
"""
# Write the record.
if not self.quiet:
print("Creating a record.")
if record is None:
record = ModelRecord()
record.name = model_name
record.wnid = f'n{wnid:08d}'
record.wcategory = wcategory
record.urls = urls
record.scale = scale
# Append asset bundle sizes.
local_path = Path.home().joinpath(
"asset_bundle_creator/Assets/NewAssetBundles").joinpath(
record.name)
for os_dir in local_path.iterdir():
if not os_dir.is_dir():
continue
asset_bundle_platform = UNITY_TO_SYSTEM[os_dir.stem]
size = os_dir.joinpath(record.name).stat().st_size
record.asset_bundle_sizes[asset_bundle_platform] = size
# Assemble a dictionary of just the data that we don't need the Editor for.
record_data = {
"name": record.name,
"urls": record.urls,
"wnid": record.wnid,
"wcategory": record.wcategory,
"scale_factor": record.scale_factor,
"do_not_use": record.do_not_use,
"do_not_use_reason": record.do_not_use_reason,
"canonical_rotation": record.canonical_rotation,
"physics_quality": -1,
"asset_bundle_sizes": record.asset_bundle_sizes
}
# Serialize the record.
record_data = json.dumps(record_data)
# Remove the last } and replace it with , to keep serializing with Unity.
record_data = record_data[:-1] + ","
record_path = self.get_assets_directory().joinpath(model_name +
".json")
record_path.write_text(record_data, encoding="utf-8")
record_call = self.unity_call[:]
record_call.extend([
"-executeMethod", "RecordCreator.WriteRecord",
"-modelname=" + model_name, "-scale=" + str(scale)
])
call(record_call)
# Test the record.
try:
json.loads(record_path.read_text(encoding="utf-8"))
except json.JSONDecodeError:
raise Exception("Failed to deserialize: " +
record_path.read_text(encoding="utf-8"))
if not self.quiet:
print("Wrote the record data to the disk.")
if write_physics:
self.write_physics_quality(
record_path=record_path,
asset_bundle_path=self.get_local_asset_bundle_path(model_name))
return record_path
def trial(self, scene: Scene, record: ModelRecord, output_path: Path,
scene_index: int) -> None:
"""
Run a trial in a scene that has been initialized.
:param scene: Data for the current scene.
:param record: The model's metadata record.
:param output_path: Write the .wav file to this path.
:param scene_index: The scene identifier.
"""
self.py_impact.reset(initial_amp=0.05)
# Initialize the scene, positioning objects, furniture, etc.
resp = self.communicate(scene.initialize_scene(self))
center = scene.get_center(self)
max_y = scene.get_max_y()
# The object's initial position.
o_x = RNG.uniform(center["x"] - 0.15, center["x"] + 0.15)
o_y = RNG.uniform(max_y - 0.5, max_y)
o_z = RNG.uniform(center["z"] - 0.15, center["z"] + 0.15)
# Physics values.
mass = self.object_info[record.name].mass + RNG.uniform(
self.object_info[record.name].mass * -0.15,
self.object_info[record.name].mass * 0.15)
static_friction = RNG.uniform(0.1, 0.3)
dynamic_friction = RNG.uniform(0.7, 0.9)
# Angles of rotation.
yaw = RNG.uniform(-30, 30)
pitch = RNG.uniform(0, 45)
roll = RNG.uniform(-45, 45)
# The force applied to the object.
force = RNG.uniform(0, 5)
# The avatar's position.
a_r = RNG.uniform(1.5, 2.2)
a_x = center["x"] + a_r
a_y = RNG.uniform(1.5, 3)
a_z = center["z"] + a_r
cam_angle_min, cam_angle_max = scene.get_camera_angles()
theta = np.radians(RNG.uniform(cam_angle_min, cam_angle_max))
a_x = np.cos(theta) * (a_x - center["x"]) - np.sin(theta) * (
a_z - center["z"]) + center["x"]
a_z = np.sin(theta) * (a_x - center["x"]) + np.cos(theta) * (
a_z - center["z"]) + center["z"]
o_id = 0
# Create the object and apply a force.
commands = [{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"position": {
"x": o_x,
"y": o_y,
"z": o_z
},
"category": record.wcategory,
"id": o_id
}, {
"$type": "set_mass",
"id": o_id,
"mass": mass
}, {
"$type": "set_physic_material",
"id": o_id,
"bounciness": self.object_info[record.name].bounciness,
"static_friction": static_friction,
"dynamic_friction": dynamic_friction
}, {
"$type": "rotate_object_by",
"angle": yaw,
"id": o_id,
"axis": "yaw",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": pitch,
"id": o_id,
"axis": "pitch",
"is_world": True
}, {
"$type": "rotate_object_by",
"angle": roll,
"id": o_id,
"axis": "roll",
"is_world": True
}, {
"$type": "apply_force_magnitude_to_object",
"magnitude": force,
"id": o_id
}, {
"$type": "send_rigidbodies",
"frequency": "always"
}, {
"$type": "send_collisions",
"enter": True,
"exit": False,
"stay": False,
"collision_types": ["obj", "env"]
}, {
"$type": "send_transforms",
"frequency": "always"
}]
# Parse bounds data to get the centroid of all objects currently in the scene.
bounds = Bounds(resp[0])
if bounds.get_num() == 0:
look_at = {"x": center["x"], "y": 0.1, "z": center["z"]}
else:
centers = []
for i in range(bounds.get_num()):
centers.append(bounds.get_center(i))
centers_x, centers_y, centers_z = zip(*centers)
centers_len = len(centers_x)
look_at = {
"x": sum(centers_x) / centers_len,
"y": sum(centers_y) / centers_len,
"z": sum(centers_z) / centers_len
}
# Add the avatar.
# Set the position at a given distance (r) from the center of the scene.
# Rotate around that position to a random angle constrained by the scene's min and max angles.
commands.extend([{
"$type": "teleport_avatar_to",
"position": {
"x": a_x,
"y": a_y,
"z": a_z
}
}, {
"$type": "look_at_position",
"position": look_at
}])
# Send the commands.
resp = self.communicate(commands)
AudioUtils.start(output_path=output_path, until=(0, 10))
# Loop until all objects are sleeping.
done = False
while not done and AudioUtils.is_recording():
commands = []
collisions, environment_collisions, rigidbodies = PyImpact.get_collisions(
resp)
# Create impact sounds from object-object collisions.
for collision in collisions:
if PyImpact.is_valid_collision(collision):
# Get the audio material and amp.
collider_id = collision.get_collider_id()
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
collidee_id = collision.get_collider_id()
collidee_material, collidee_amp = self._get_object_info(
collidee_id, Scene.OBJECT_IDS, record.name)
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collidee_id,
target_amp=collidee_amp,
target_mat=collidee_material.name,
other_id=collider_id,
other_mat=collider_material.name,
other_amp=collider_amp,
play_audio_data=False)
commands.append(impact_sound_command)
# Create impact sounds from object-environment collisions.
for collision in environment_collisions:
collider_id = collision.get_object_id()
if self._get_velocity(rigidbodies, collider_id) > 0:
collider_material, collider_amp = self._get_object_info(
collider_id, Scene.OBJECT_IDS, record.name)
surface_material = scene.get_surface_material()
impact_sound_command = self.py_impact.get_impact_sound_command(
collision=collision,
rigidbodies=rigidbodies,
target_id=collider_id,
target_amp=collider_amp,
target_mat=collider_material.name,
other_id=-1,
other_amp=0.01,
other_mat=surface_material.name,
play_audio_data=False)
commands.append(impact_sound_command)
# If there were no collisions, check for movement. If nothing is moving, the trial is done.
if len(commands) == 0:
transforms = AudioDataset._get_transforms(resp)
done = True
for i in range(rigidbodies.get_num()):
if self._is_moving(rigidbodies.get_id(i), transforms,
rigidbodies):
done = False
break
# Continue the trial.
if not done:
resp = self.communicate(commands)
# Stop listening for anything except audio data..
resp = self.communicate([{
"$type": "send_rigidbodies",
"frequency": "never"
}, {
"$type": "send_transforms",
"frequency": "never"
}, {
"$type": "send_collisions",
"enter": False,
"exit": False,
"stay": False,
"collision_types": []
}, {
"$type": "send_audio_sources",
"frequency": "always"
}])
# Wait for the audio to finish.
done = False
while not done and AudioUtils.is_recording():
done = True
for r in resp[:-1]:
if OutputData.get_data_type_id(r) == "audi":
audio_sources = AudioSources(r)
for i in range(audio_sources.get_num()):
if audio_sources.get_is_playing(i):
done = False
if not done:
resp = self.communicate([])
# Cleanup.
commands = [{
"$type": "send_audio_sources",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": o_id
}]
for scene_object_id in Scene.OBJECT_IDS:
commands.append({"$type": "destroy_object", "id": scene_object_id})
self.communicate(commands)
# Insert the trial's values into the database.
self.db_c.execute(
"INSERT INTO sound20k VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
(output_path.name, scene_index, a_x, a_y, a_z, o_x, o_y, o_z, mass,
static_friction, dynamic_friction, yaw, pitch, roll, force))
self.conn.commit()
"object_id": object_id,
"use_centroid": True}
])
grades = [1 - (float(h) / float(i)) for h, i in zip(pink_pass, id_pass)]
physics_quality = float(sum(grades)) / len(grades)
print("Physics quality: " + str(physics_quality))
# Kill the build.
self.kill_build()
return physics_quality
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--record_path", type=str, help="The path to the temporary record file")
parser.add_argument("--asset_bundle_path", type=str, help="The path to the local asset bundle.")
args = parser.parse_args()
# Get the physics quality.
c = PhysicsQualityWriter(record_path=args.record_path,
asset_bundle_path=args.asset_bundle_path)
physics_quality = c.get_physics_quality()
# Update the record.
record_path = Path(args.record_path)
record = ModelRecord(json.loads(record_path.read_text(encoding="utf-8")))
record.physics_quality = physics_quality
record_path.write_text(json.dumps(record.__dict__), encoding="utf-8")
def process_model(self, record: ModelRecord, a: str, envs: list,
train_count: int, val_count: int, root_dir: str,
wnid: str) -> float:
"""
Capture images of a model.
:param record: The model record.
:param a: The ID of the avatar.
:param envs: All environment data.
:param train_count: Number of train images.
:param val_count: Number of val images.
:param root_dir: The root directory for saving images.
:param wnid: The wnid of the record.
:return The time elapsed.
"""
image_count = 0
# Get the filename index. If we shouldn't overwrite any images, start after the last image.
if self.no_overwrite:
# Check if any images exist.
wnid_dir = Path(root_dir).joinpath(f"train/{wnid}")
if wnid_dir.exists():
max_file_index = -1
for image in wnid_dir.iterdir():
if not image.is_file() or image.suffix != ".jpg" \
or not image.stem.startswith("img_") or image.stem[4:-5] != record.name:
continue
image_index = int(image.stem[-4:])
if image_index > max_file_index:
max_file_index = image_index
file_index = max_file_index + 1
else:
file_index = 0
else:
file_index = 0
image_positions = []
o_id = self.get_unique_id()
s = TDWUtils.get_unit_scale(record)
# Add the object.
# Set the screen size to 32x32 (to make the build run faster; we only need the average grayscale values).
# Toggle off pass masks.
# Set render quality to minimal.
# Scale the object to "unit size".
self.communicate([{
"$type": "add_object",
"name": record.name,
"url": record.get_url(),
"scale_factor": record.scale_factor,
"category": record.wcategory,
"id": o_id
}, {
"$type": "set_screen_size",
"height": 32,
"width": 32
}, {
"$type": "set_pass_masks",
"avatar_id": a,
"pass_masks": []
}, {
"$type": "set_render_quality",
"render_quality": 0
}, {
"$type": "scale_object",
"id": o_id,
"scale_factor": {
"x": s,
"y": s,
"z": s
}
}])
# The index in the HDRI records array.
hdri_index = 0
# The number of iterations on this skybox so far.
skybox_count = 0
if self.skyboxes:
# The number of iterations per skybox for this model.
its_per_skybox = round(
(train_count + val_count) / len(self.skyboxes))
# Set the first skybox.
hdri_index, skybox_count, command = self.set_skybox(
self.skyboxes, its_per_skybox, hdri_index, skybox_count)
self.communicate(command)
else:
its_per_skybox = 0
while len(image_positions) < train_count + val_count:
e = RNG.choice(envs)
# Get the real grayscale.
g_r, d, a_p, o_p, o_rot, cam_rot = self.get_real_grayscale(
o_id, a, e)
if g_r > 0:
# Get the optimal grayscale.
g_o = self.get_optimal_grayscale(o_id, a, o_p, a_p)
if g_o > 0 and g_r / g_o > self.grayscale_threshold:
# Cache the position.
image_positions.append(
ImagePosition(a_p, cam_rot, o_p, o_rot))
# Send images.
# Set the screen size.
# Set render quality to maximum.
commands = [{
"$type": "send_images",
"frequency": "always"
}, {
"$type": "set_pass_masks",
"avatar_id": a,
"pass_masks": ["_img", "_id"] if self.id_pass else ["_img"]
}, {
"$type": "set_screen_size",
"height": self.screen_size,
"width": self.screen_size
}, {
"$type": "set_render_quality",
"render_quality": 5
}]
# Hide the object maybe.
if not self.show_objects:
commands.append({"$type": "hide_object", "id": o_id})
self.communicate(commands)
t0 = time()
# Generate images from the cached spatial data.
train = 0
for p in image_positions:
# Teleport the avatar.
# Rotate the avatar's camera.
# Teleport the object.
# Rotate the object.
# Get the response.
commands = [{
"$type": "teleport_avatar_to",
"avatar_id": a,
"position": p.avatar_position
}, {
"$type": "rotate_sensor_container_to",
"avatar_id": a,
"rotation": p.camera_rotation
}, {
"$type": "teleport_object",
"id": o_id,
"position": p.object_position
}, {
"$type": "rotate_object_to",
"id": o_id,
"rotation": p.object_rotation
}]
# Set the visual materials.
if self.materials is not None:
if record.name not in self.substructures:
self.substructures.update(
{record.name: record.substructure})
for sub_object in self.substructures[record.name]:
for i in range(
len(self.substructures[record.name][
sub_object["name"]])):
material_name = self.materials[RNG.randint(
0, len(self.materials))].name
commands.extend([
self.get_add_material(material_name), {
"$type": "set_visual_material",
"id": o_id,
"material_name": material_name,
"object_name": sub_object["name"],
"material_index": i
}
])
# Maybe set a new skybox.
# Rotate the skybox.
if self.skyboxes:
hdri_index, skybox_count, command = self.set_skybox(
self.skyboxes, its_per_skybox, hdri_index, skybox_count)
if command:
commands.append(command)
commands.append({
"$type": "rotate_hdri_skybox_by",
"angle": RNG.uniform(0, 360)
})
resp = self.communicate(commands)
train += 1
# Create a thread to save the image.
t = Thread(target=self.save_image,
args=(resp, record, file_index, root_dir, wnid, train,
train_count))
t.daemon = True
t.start()
file_index += 1
image_count += 1
t1 = time()
# Stop sending images.
# Destroy the object.
# Unload asset bundles.
self.communicate([{
"$type": "send_images",
"frequency": "never"
}, {
"$type": "destroy_object",
"id": o_id
}, {
"$type": "unload_asset_bundles"
}])
return t1 - t0
# Create asset bundles.
a = AssetBundleCreator()
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