def run(self): self.start() # Create an empty room. self.communicate(TDWUtils.create_empty_room(12, 12)) material_librarian = MaterialLibrarian() # Fetch a random wood material record. wood_records = material_librarian.get_all_materials_of_type("Wood") wood_record: MaterialRecord = choice(wood_records) # Fetch a random concrete record material. concrete_records = material_librarian.get_all_materials_of_type( "Concrete") concrete_record: MaterialRecord = choice(concrete_records) # Set the floor material. # Set the floor material texture scale. # Set the walls material. # Set the wall material texture scale. self.communicate([ self.get_add_material(wood_record.name), self.get_add_material(concrete_record.name), { "$type": "set_proc_gen_floor_material", "name": wood_record.name }, { "$type": "set_proc_gen_floor_texture_scale", "scale": { "x": 8, "y": 8 } }, { "$type": "set_proc_gen_walls_material", "name": concrete_record.name }, { "$type": "set_proc_gen_walls_texture_scale", "scale": { "x": 2, "y": 8 } } ]) # Create a random table. model_librarian = ModelLibrarian() wnids = model_librarian.get_model_wnids_and_wcategories() table_wnid = [wnid for wnid in wnids if wnids[wnid] == "table"][0] table_records = model_librarian.get_all_models_in_wnid(table_wnid) # Filter out unusable tables. table_records = [t for t in table_records if not t.do_not_use] table_record: ModelRecord = choice(table_records) # Add the table. table_id = self.add_object(table_record.name) # Set a random visual material for each sub-object in the model. commands = list() for sub_object in table_record.substructure: for i in range(len(sub_object["materials"])): material_record: MaterialRecord = choice( material_librarian.records) commands.extend([ self.get_add_material(material_record.name), { "$type": "set_visual_material", "id": table_id, "material_name": material_record.name, "object_name": sub_object["name"], "material_index": i } ]) self.communicate(commands) # Create the avatar. self.communicate( TDWUtils.create_avatar(position={ "x": 4, "y": 3, "z": 2 }, look_at=TDWUtils.VECTOR3_ZERO))
class SingleObject(Controller): def __init__(self, port=1071, visual_material_swapping=False, new=False, screen_size=256, output_size=256, hdri=False, show_objects=True, clamp_rotation=False, max_height=1.0, grayscale_threshold=0.5, less_dark=False, id_pass=False, no_overwrite=False, do_zip=False, train=1300000, val=50000, library="models_full.json"): """ :param port: The port used to connect to the build. :param visual_material_swapping: If true, set random visual materials per frame. :param new: If true, clear the list of models that have already been used. :param screen_size: The screen size of the build. :param output_size: The size of the output images. :param hdri: If true, use a random HDRI skybox per frame. :param show_objects: If true, show objects. :param clamp_rotation: If true, clamp the rotation to +/- 30 degrees around each axis. :param max_height: The percentage of the environment height that is the ceiling for the avatar and object. Must be between 0 and 1. :param grayscale_threshold: The grayscale threshold. Higher value = slower FPS, better composition. Must be between 0 and 1. :param less_dark: If true, there will be more daylight exterior skyboxes (requires hdri == True) :param id_pass: If true, send the _id pass. :param no_overwrite: If true, don't overwrite images. :param do_zip: If true, zip the directory at the end. :param train: Number of train images. :param val: Number of val images. :param library: The path to the library records file. """ self.screen_size = screen_size self.output_size = output_size self.show_objects = show_objects self.clamp_rotation = clamp_rotation self.max_height = max_height self.grayscale_threshold = grayscale_threshold self.id_pass = id_pass self.no_overwrite = no_overwrite self.do_zip = do_zip self.train = train self.val = val assert 0 < max_height <= 1.0, f"Invalid max height: {max_height}" assert 0 < grayscale_threshold <= 1.0, f"Invalid grayscale threshold: {grayscale_threshold}" self.less_dark = less_dark self.substructures: Dict[str, List[dict]] = {} self.new = new super().__init__(port=port) self.model_librarian = ModelLibrarian(library=library) self.material_librarian = MaterialLibrarian("materials_high.json") self.hdri_skybox_librarian = HDRISkyboxLibrarian() # Get material records. if visual_material_swapping: self.materials = self.material_librarian.records else: self.materials = None # Get skybox records. if hdri: self.skyboxes = self.hdri_skybox_librarian.records # Prefer exterior daytime skyboxes by adding them multiple times to the list. if self.less_dark: temp = self.skyboxes[:] for skybox in temp: if skybox.location != "interior" and skybox.sun_elevation >= 145: self.skyboxes.append(skybox) else: self.skyboxes = None def initialize_scene(self, scene_command, a="a") -> list: """ Initialize the scene. :param scene_command: The command to load the scene. :param a: The avatar ID. :return: The Environments data of the scene. """ # Initialize the scene. # Add the avatar. commands = [ scene_command, { "$type": "create_avatar", "type": "A_Img_Caps_Kinematic", "id": a, "envs": [0] } ] # Disable physics. # Enable jpgs. # Set FOV. # Set clipping planes. # Set AA. # Set aperture. # Disable vignette. commands.extend([{ "$type": "simulate_physics", "value": False }, { "$type": "set_img_pass_encoding", "value": False }, { '$type': 'set_field_of_view', 'avatar_id': 'a', 'field_of_view': 60 }, { '$type': 'set_camera_clipping_planes', 'avatar_id': 'a', 'far': 160, 'near': 0.01 }, { "$type": "set_anti_aliasing", "avatar_id": "a", "mode": "subpixel" }, { "$type": "set_aperture", "aperture": 70 }, { '$type': 'set_vignette', 'enabled': False }]) # If we're using HDRI skyboxes, send additional favorable post-process commands. if self.skyboxes is not None: commands.extend([{ "$type": "set_post_exposure", "post_exposure": 0.6 }, { "$type": "set_contrast", "contrast": -20 }, { "$type": "set_saturation", "saturation": 10 }, { "$type": "set_screen_space_reflections", "enabled": False }, { "$type": "set_shadow_strength", "strength": 1.0 }]) # Send the commands. self.communicate(commands) # Get the environments data. env_data = Environments( self.communicate({ "$type": "send_environments", "frequency": "once" })[0]) envs = [] for i in range(env_data.get_num()): envs.append(Environment(env_data, i)) return envs def generate_metadata(self, dataset_dir: str, scene_name: str) -> None: """ Generate a metadata file for this dataset. :param dataset_dir: The dataset directory for images. :param scene_name: The scene name. """ root_dir = f"{dataset_dir}/images/" if not os.path.exists(root_dir): os.makedirs(root_dir) data = { "dataset": dataset_dir, "scene": scene_name, "train": self.train, "val": self.val, "visual_material_swapping": self.materials is not None, "hdri": self.skyboxes is not None, "screen_size": self.screen_size, "output_size": self.output_size, "clamp_rotation": self.clamp_rotation, "show_objects": self.show_objects, "max_height": self.max_height, "grayscale_threshold": self.grayscale_threshold, "less_dark": self.less_dark, "multi_scene": self.no_overwrite, "start": datetime.now().strftime("%H:%M %d.%m.%y") } with open(os.path.join(root_dir, "metadata.txt"), "wt") as f: json.dump(data, f, sort_keys=True, indent=4) def run(self, dataset_dir: str, scene_name: str) -> None: """ Generate the dataset. :param dataset_dir: The dataset directory for images. :param scene_name: The scene name. """ # Create the metadata file. self.generate_metadata(dataset_dir, scene_name=scene_name) # The root directory of the output. root_dir = f"{dataset_dir}/images/" # The avatar ID. a = "a" # Initialize the scene. envs = self.initialize_scene(self.get_add_scene(scene_name)) # Fetch the WordNet IDs. wnids = self.model_librarian.get_model_wnids() # Remove any wnids that don't have valid models. wnids = [ w for w in wnids if len([ r for r in self.model_librarian.get_all_models_in_wnid(w) if not r.do_not_use ]) > 0 ] # Set the number of train and val images per wnid. num_train = self.train / len(wnids) num_val = self.val / len(wnids) # Create the progress bar. pbar = tqdm(total=len(wnids)) # If this is a new dataset, remove the previous list of completed models. done_models_filename = "processed_records.txt" if self.new and os.path.exists(done_models_filename): os.remove(done_models_filename) # Get a list of models that have already been processed. processed_model_names = [] if os.path.exists(done_models_filename): with open(done_models_filename, "rt") as f: txt = f.read() processed_model_names = txt.split("\n") # Iterate through each wnid. for w, q in zip(wnids, range(len(wnids))): # Update the progress bar. pbar.set_description(w) # Get all valid models in the wnid. records = self.model_librarian.get_all_models_in_wnid(w) records = [r for r in records if not r.do_not_use] # Get the train and val counts. train_count = [ len(a) for a in np.array_split(np.arange(num_train), len(records)) ][0] val_count = [ len(a) for a in np.array_split(np.arange(num_val), len(records)) ][0] # Process each record. fps = "nan" for record, i in zip(records, range(len(records))): # Set the progress bar description to the wnid and FPS. pbar.set_description( f"record {i + 1}/{len(records)}, FPS {fps}") # Skip models that have already been processed. if record.name in processed_model_names: continue # Create all of the images for this model. dt = self.process_model(record, a, envs, train_count, val_count, root_dir, w) fps = round((train_count + val_count) / dt) # Mark this record as processed. with open(done_models_filename, "at") as f: f.write(f"\n{record.name}") pbar.update(1) pbar.close() # Add the end time to the metadata file. with open(os.path.join(root_dir, "metadata.txt"), "rt") as f: data = json.load(f) end_time = datetime.now().strftime("%H:%M %d.%m.%y") if "end" in data: data["end"] = end_time else: data.update({"end": end_time}) with open(os.path.join(root_dir, "metadata.txt"), "wt") as f: json.dump(data, f, sort_keys=True, indent=4) # Terminate the build. if not self.no_overwrite: self.communicate({"$type": "terminate"}) # Zip up the images. if self.do_zip: zip_dir = Path(dataset_dir) SingleObject.zip_images(zip_dir) def set_skybox(self, records: List[HDRISkyboxRecord], its_per_skybox: int, hdri_index: int, skybox_count: int) -> (int, int, dict): """ If it's time, set a new skybox. :param records: All HDRI records. :param its_per_skybox: Iterations per skybox. :param hdri_index: The index in the records list. :param skybox_count: The number of images of this model with this skybox. :return: (hdri_index, skybox_count, command used to set the skybox) """ # Set a new skybox. if skybox_count == 0: command = self.get_add_hdri_skybox(records[hdri_index].name) # It's not time yet to set a new skybox. Don't send a command. else: command = None skybox_count += 1 if skybox_count >= its_per_skybox: hdri_index += 1 if hdri_index >= len(records): hdri_index = 0 skybox_count = 0 return hdri_index, skybox_count, command 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 def save_image(self, resp, record: ModelRecord, image_count: int, root_dir: str, wnid: str, train: int, train_count: int) -> None: """ Save an image. :param resp: The raw response data. :param record: The model record. :param image_count: The image count. :param root_dir: The root directory. :param wnid: The wnid. :param train: Number of train images so far. :param train_count: Total number of train images to generate. """ # Get the directory. directory = Path(root_dir).joinpath("train" if train < train_count else "val").joinpath(wnid).resolve() if not os.path.exists(directory): # Try to make the directories. Due to threading, they might already be made. try: os.makedirs(directory) except: pass # Save the image. filename = f"{record.name}_{image_count:04d}" # Save the image without resizing. if self.screen_size == self.output_size: TDWUtils.save_images(Images(resp[0]), filename, output_directory=directory) # Resize the image and save it. else: TDWUtils.save_images(Images(resp[0]), filename, output_directory=directory, resize_to=(self.output_size, self.output_size)) def get_optimal_grayscale(self, o_id: int, a_id: str, o_p: dict, a_p: dict) -> float: """ Get the "optimal" grayscale value of the object if there wasn't any occlusion. :param o_id: The ID of the object. :param a_id: The ID of the avatar. :param o_p: The position of the object. :param a_p: The position of the avatar. :return: The grayscale value. """ # Teleport the object into the sky. # Teleport the avatar into the sky. # Return the grayscale value. return IdPassGrayscale( self.communicate([{ "$type": "teleport_object", "id": o_id, "position": { "x": o_p["x"], "y": o_p["y"] + 500, "z": o_p["z"] } }, { "$type": "teleport_avatar_to", "avatar_id": a_id, "position": { "x": a_p["x"], "y": a_p["y"] + 500, "z": a_p["z"] } }, { "$type": "send_id_pass_grayscale", "frequency": "once" }])[0]).get_grayscale() def get_real_grayscale( self, o_id: int, a_id: str, e: Environment) -> (float, float, dict, dict, dict, dict): """ Get the "real" grayscale value of an image we hope to capture. :param o_id: The ID of the object. :param a_id: The ID of the avatar. :param e: The environment. :return: (grayscale, distance, avatar_position, object_position, object_rotation, avatar_rotation) """ # Get a random position for the avatar. a_p = np.array([ RNG.uniform(-e.w / 2, e.w / 2) + e.x, RNG.uniform(0.4, e.h * self.max_height), RNG.uniform(-e.l / 2, e.l / 2) + e.z ]) # Get a random distance from the avatar. d = RNG.uniform(0.8, 3) # Get a random position for the object constrained to the environment bounds. o_p = SingleObject.sample_spherical() * d # Clamp the y value to positive. o_p[1] = abs(o_p[1]) o_p = a_p + o_p # Clamp the y value of the object. if o_p[1] > e.h * self.max_height: o_p[1] = e.h * self.max_height # Convert the avatar's position to a Vector3. a_p = TDWUtils.array_to_vector3(a_p) # Set random camera rotations. yaw = RNG.uniform(-15, 15) pitch = RNG.uniform(-15, 15) # Convert the object position to a Vector3. o_p = TDWUtils.array_to_vector3(o_p) # Add rotation commands. # If we're clamping the rotation, rotate the object within +/- 30 degrees on each axis. if self.clamp_rotation: o_rot = None commands = [{ "$type": "rotate_object_to", "id": o_id, "rotation": { "x": 0, "y": 0, "z": 0, "w": 0 } }, { "$type": "rotate_object_by", "id": o_id, "angle": RNG.uniform(-30, 30), "axis": "pitch" }, { "$type": "rotate_object_by", "id": o_id, "angle": RNG.uniform(-30, 30), "axis": "yaw" }, { "$type": "rotate_object_by", "id": o_id, "angle": RNG.uniform(-30, 30), "axis": "roll" }] # Set a totally random rotation. else: o_rot = { "x": RNG.uniform(-360, 360), "y": RNG.uniform(-360, 360), "z": RNG.uniform(-360, 360), "w": RNG.uniform(-360, 360), } commands = [{ "$type": "rotate_object_to", "id": o_id, "rotation": o_rot }] # After rotating the object: # 1. Teleport the object. # 2. Teleport the avatar. # 3. Look at the object. # 4. Perturb the camera slightly. # 5. Send grayscale data and image sensor data. commands.extend([{ "$type": "teleport_object", "id": o_id, "position": o_p }, { "$type": "teleport_avatar_to", "avatar_id": a_id, "position": a_p }, { "$type": "look_at", "avatar_id": a_id, "object_id": o_id, "use_centroid": True }, { "$type": "rotate_sensor_container_by", "angle": pitch, "avatar_id": "a", "axis": "pitch" }, { "$type": "rotate_sensor_container_by", "angle": yaw, "avatar_id": "a", "axis": "yaw" }, { "$type": "send_id_pass_grayscale", "frequency": "once" }, { "$type": "send_image_sensors", "frequency": "once" }]) resp = self.communicate(commands) # Parse the output data: # 1. The grayscale value of the image. # 2. The camera rotation. grayscale = 0 cam_rot = None for r in resp[:-1]: r_id = OutputData.get_data_type_id(r) if r_id == "idgs": grayscale = IdPassGrayscale(r).get_grayscale() elif r_id == "imse": cam_rot = ImageSensors(r).get_sensor_rotation(0) cam_rot = { "x": cam_rot[0], "y": cam_rot[1], "z": cam_rot[2], "w": cam_rot[3] } # If we clamped the rotation of the object, we need to know its quaternion. if self.clamp_rotation: resp = self.communicate({ "$type": "send_transforms", "frequency": "once", "ids": [o_id] }) t = Transforms(resp[0]) o_rot = t.get_rotation(0) o_rot = { "x": o_rot[0], "y": o_rot[1], "z": o_rot[2], "w": o_rot[3], } return grayscale, d, a_p, o_p, o_rot, cam_rot @staticmethod def sample_spherical(npoints=1, ndim=3) -> np.array: vec = RNG.randn(ndim, npoints) vec /= np.linalg.norm(vec, axis=0) return np.array([vec[0][0], vec[1][0], vec[2][0]]) @staticmethod def zip_images(zip_dir: Path) -> None: """ Zip up the images. :param zip_dir: The zip directory. """ if not zip_dir.exists(): zip_dir.mkdir() # Use a random token to avoid overwriting zip files. token = token_urlsafe(4) zip_path = zip_dir.joinpath(f"images_{token}.zip") images_directory = str(zip_dir.joinpath("images").resolve()) # Create the zip file. If it is made, remove the original directory. zip_call = f'C:/Program Files/7-Zip/7z.exe a -r {str(zip_path.resolve())} {images_directory} -sdel' call(zip_call)