(self, idx, resolution, rng)
| 63 | return depthmap |
| 64 | |
| 65 | def _get_views(self, idx, resolution, rng): |
| 66 | # choose a scene |
| 67 | obj, instance = self.scene_list[idx // len(self.combinations)] |
| 68 | image_pool = self.scenes[obj, instance] |
| 69 | im1_idx, im2_idx = self.combinations[idx % len(self.combinations)] |
| 70 | |
| 71 | # add a bit of randomness |
| 72 | last = len(image_pool) - 1 |
| 73 | |
| 74 | if resolution not in self.invalidate[obj, instance]: # flag invalid images |
| 75 | self.invalidate[obj, instance][resolution] = [False for _ in range(len(image_pool))] |
| 76 | |
| 77 | # decide now if we mask the bg |
| 78 | mask_bg = (self.mask_bg == True) or (self.mask_bg == 'rand' and rng.choice(2)) |
| 79 | |
| 80 | views = [] |
| 81 | imgs_idxs = [max(0, min(im_idx + rng.integers(-4, 5), last)) for im_idx in [im2_idx, im1_idx]] |
| 82 | imgs_idxs = deque(imgs_idxs) |
| 83 | while len(imgs_idxs) > 0: # some images (few) have zero depth |
| 84 | im_idx = imgs_idxs.pop() |
| 85 | |
| 86 | if self.invalidate[obj, instance][resolution][im_idx]: |
| 87 | # search for a valid image |
| 88 | random_direction = 2 * rng.choice(2) - 1 |
| 89 | for offset in range(1, len(image_pool)): |
| 90 | tentative_im_idx = (im_idx + (random_direction * offset)) % len(image_pool) |
| 91 | if not self.invalidate[obj, instance][resolution][tentative_im_idx]: |
| 92 | im_idx = tentative_im_idx |
| 93 | break |
| 94 | |
| 95 | view_idx = image_pool[im_idx] |
| 96 | |
| 97 | impath = self._get_impath(obj, instance, view_idx) |
| 98 | depthpath = self._get_depthpath(obj, instance, view_idx) |
| 99 | |
| 100 | # load camera params |
| 101 | metadata_path = self._get_metadatapath(obj, instance, view_idx) |
| 102 | input_metadata = np.load(metadata_path) |
| 103 | camera_pose = input_metadata['camera_pose'].astype(np.float32) |
| 104 | intrinsics = input_metadata['camera_intrinsics'].astype(np.float32) |
| 105 | |
| 106 | # load image and depth |
| 107 | rgb_image = imread_cv2(impath) |
| 108 | depthmap = self._read_depthmap(depthpath, input_metadata) |
| 109 | |
| 110 | if mask_bg: |
| 111 | # load object mask |
| 112 | maskpath = self._get_maskpath(obj, instance, view_idx) |
| 113 | maskmap = imread_cv2(maskpath, cv2.IMREAD_UNCHANGED).astype(np.float32) |
| 114 | maskmap = (maskmap / 255.0) > 0.1 |
| 115 | |
| 116 | # update the depthmap with mask |
| 117 | depthmap *= maskmap |
| 118 | |
| 119 | rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary( |
| 120 | rgb_image, depthmap, intrinsics, resolution, rng=rng, info=impath) |
| 121 | |
| 122 | num_valid = (depthmap > 0.0).sum() |
nothing calls this directly
no test coverage detected