diff --git a/objathor/dataset/generate_holodeck_features.py b/objathor/dataset/generate_holodeck_features.py index dd102eb..2a50c91 100644 --- a/objathor/dataset/generate_holodeck_features.py +++ b/objathor/dataset/generate_holodeck_features.py @@ -56,28 +56,24 @@ def __init__( if self.image_preprocessor is not None: assert len(img_angles) > 0, "At least one angle must be provided." - print("Globbing to find all thor renders...") - render_paths = glob.glob( - os.path.join( - self.asset_dir, f"*/thor_renders/0_1_0_{img_angles[0]:.1f}.*" - ) - ) + for uid in tqdm.tqdm( + self.uids, + desc="Checking renders exist for all successfully processed objects.", + ): + for angle in img_angles: + rp = os.path.join( + self.asset_dir, + uid, + "thor_renders", + f"0_1_0_{img_angles[0]:.1f}.jpg", + ) + if not os.path.exists(rp): + rp = rp.replace(".jpg", ".png") - for rp in tqdm.tqdm(render_paths): - for angle in img_angles[1:]: assert os.path.exists( - rp.replace(f"0_1_0_{img_angles[0]:.1f}", f"0_1_0_{angle:.1f}") + rp ), f"Missing render for {os.path.dirname(rp)} at angle {angle}." - render_uids = set( - os.path.basename(os.path.dirname(os.path.dirname(p))) - for p in render_paths - ) - - assert ( - len(set(self.uids) - render_uids) == 0 - ), f"Some objects with annotations are missing renders: {set(self.uids) - render_uids}." - def __len__(self) -> int: return len(self.uids)