def __getitem__(self, idx):
if idx < 0 or idx >= self.__len__():
return None
sparse_img = Image.open(str(self.sparse_depth_paths[idx]))
gt_img = Image.open(str(self.gt_paths[idx]))
# convert to meters
# TODO float16 or float32
sparse_img = np.array(sparse_img, dtype=np.float32) / self.norm_factor
gt_img = np.array(gt_img, dtype=np.float32) / self.norm_factor
sparse_img = self._crop_img(sparse_img)
gt_img = self._crop_img(gt_img)
# convert to torch tensor
sparse_img = utils.to_torch(sparse_img[None])
gt_img = utils.to_torch(gt_img[None])
if self.load_rgb:
rgb_img = self._read_rgb_image(str(self.rgb_paths[idx]))
result_dict = {
keys.SPARSE_DEPTH: sparse_img,
keys.ANNOTATED_DEPTH: gt_img,
}
if self.load_rgb:
result_dict[keys.ALIGNED_RGB] = rgb_img
return result_dict
def get_frame_space_scaling_homography():
src_4pts = utils.to_torch(utils.FULL_CANON4PTS_NP())
dest_4pts = utils.to_torch(
np.array([[0, 0], [0, self.frame_h],
[self.frame_w, self.frame_h], [self.frame_w, 0]],
dtype=np.float32))
scaling_transformation = utils.get_perspective_transform(
src_4pts[None], dest_4pts[None])
scaling_transformation = utils.to_numpy(scaling_transformation[0])
return scaling_transformation
def get_default_canon4pts(batch_size, canon4pts_type: str):
if canon4pts_type == 'lower':
lower_canon4pts = utils.LOWER_CANON4PTS_NP()
lower_canon4pts = np.tile(lower_canon4pts, (batch_size, 1, 1))
lower_canon4pts = utils.to_torch(lower_canon4pts)
return lower_canon4pts
elif canon4pts_type == 'full':
full_canon4pts = utils.FULL_CANON4PTS_NP()
full_canon4pts = np.tile(full_canon4pts, (batch_size, 1, 1))
full_canon4pts = utils.to_torch(full_canon4pts)
return full_canon4pts
else:
raise ValueError('unknown canon4pts type')
def np_img_to_torch_img(np_img):
'''convert a numpy image to bilinear-samplable torch image
numpy use Height x Width x Channels
torch use Channels x Height x Width
Arguments:
np_img {[type]} -- [description]
'''
if len(np_img.shape) == 4 and (np_img.shape[3] == 3 or np_img.shape[3] == 1):
return utils.to_torch(np.transpose(np_img, (0, 3, 1, 2)))
if len(np_img.shape) == 3 and (np_img.shape[2] == 3 or np_img.shape[2] == 1):
return utils.to_torch(np.transpose(np_img, (2, 0, 1)))
elif len(np_img.shape) == 2:
return utils.to_torch(np_img)
else:
raise ValueError('cannot process this image')
def get_warped_tmp_by_id(self, data_id):
homo_mat = self.get_homography_by_id(data_id)
warped_tmp = warp.warp_image(self.template_torch,
utils.to_torch(homo_mat),
out_shape=(self.frame_h, self.frame_w))
warped_tmp = utils.torch_img_to_np_img(warped_tmp[0])
return warped_tmp
def get_four_corners(homo_mat, canon4pts=None):
'''
calculate the 4 corners after transformation, from frame to template
assuming the original 4 corners of the frame are [+-0.5, +-0.5]
note: this function supports batch processing
Arguments:
homo_mat {[type]} -- [homography, shape: (B, 3, 3) or (3, 3)]
Return:
xy_warped -- torch.Size([B, 2, 4])
'''
# append ones for homogeneous coordinates
if homo_mat.shape == (3, 3):
homo_mat = homo_mat[None]
assert homo_mat.shape[1:] == (3, 3)
if canon4pts is None:
canon4pts = utils.to_torch(utils.FULL_CANON4PTS_NP())
assert canon4pts.shape == (4, 2)
x, y = canon4pts[:, 0], canon4pts[:, 1]
xy = torch.stack([x, y, torch.ones_like(x)])
# warp points to model coordinates
xy_warped = torch.matmul(homo_mat, xy) # H.bmm(xy)
xy_warped, z_warped = xy_warped.split(2, dim=1)
xy_warped = xy_warped / (z_warped + 1e-8)
return xy_warped
def get_image_by_index(self, index):
img = self.raw_data.root.frames[index]
img = img[..., [2, 1, 0]]
img = img / 255.0
img = utils.to_torch(img).permute(2, 0, 1)
if self.opt.need_single_image_normalization:
img = image_utils.normalize_single_image(img)
return img
def _read_rgb_image(self, rgb_path: str):
rgb_img = Image.open(rgb_path)
# normalize the color
rgb_img = (np.array(rgb_img, dtype=np.float32) / 255.0) * 2.0 - 1.0
rgb_img = self._crop_img(rgb_img)
rgb_img = utils.to_torch(np.transpose(rgb_img, (2, 0, 1)))
rgb_img = image_utils.normalize_single_image(rgb_img)
# exec(utils.TEST_EMBEDDING)
return rgb_img
def get_homography_by_index(self, index):
homography = self.raw_data.root.homographies[index]
# if hasattr(self, 'raw_data.root'):
# homography = self.raw_data.root.homographies[index]
# else:
# homography = self.homographies[index]
homography = utils.to_torch(homography)
homography = homography / homography[2:3, 2:3]
return homography
def squash_image(image):
'''squash image to [0, 1]
'''
min_val = image.min()
max_val = image.max()
if isinstance(image, np.ndarray):
image = np.interp(image, (min_val, max_val), (0, 1))
return image
elif isinstance(image, torch.Tensor):
image_np = utils.to_numpy(image)
image_np = np.interp(image_np, (min_val, max_val), (0, 1))
image_torch = utils.to_torch(image_np)
return image_torch
else:
raise ValueError('unsupported data type: {0}'.format(type(image)))
def load_template(self):
self.template_np = imageio.imread(self.template_path,
pilmode='RGB') / 255.0
self.template_torch = utils.to_torch(
imageio.imread(self.template_path, pilmode='RGB') / 255.0).permute(
2, 0, 1)
self.template_torch = torch.mean(self.template_torch, 0, keepdim=True)
x_coord, y_coord = torch.meshgrid([
torch.linspace(-1, 1, steps=self.template_torch.shape[-2]),
torch.linspace(-1, 1, steps=self.template_torch.shape[-1])
])
x_coord = x_coord.to(self.template_torch.device)[None]
y_coord = y_coord.to(self.template_torch.device)[None]
self.template_torch = torch.cat(
[self.template_torch, x_coord, y_coord], dim=0)
def __getitem__(self, idx):
sparse_img = Image.open(str(self.sparse_depth_paths[idx]))
gt_img = Image.open(str(self.gt_paths[idx]))
sparse_nn_depth, sparse_confidence_map = self._read_nn_filled_image(
str(self.sparse_nn_depth_paths[idx]))
gt_nn_depth, gt_confidence_map = self._read_nn_filled_image(
str(self.gt_nn_paths[idx]))
# convert to meters
sparse_img = np.array(sparse_img, dtype=np.float32) / self.norm_factor
gt_img = np.array(gt_img, dtype=np.float32) / self.norm_factor
sparse_img = self._crop_img(sparse_img)
gt_img = self._crop_img(gt_img)
# exec(utils.TEST_EMBEDDING)
# sparse_nn_depth, sparse_confidence_map = image_utils.gen_voronoi_from_sparse(self.opt, sparse_img)
# gt_nn_depth, gt_confidence_map = image_utils.gen_voronoi_from_sparse(self.opt, gt_img)
# convert to torch tensor
sparse_img = utils.to_torch(sparse_img[None])
gt_img = utils.to_torch(gt_img[None])
sparse_nn_depth = utils.to_torch(sparse_nn_depth[None])
gt_nn_depth = utils.to_torch(gt_nn_depth[None])
# normalize confidence map
sparse_confidence_map = (utils.to_torch(sparse_confidence_map[None]) /
255.0) * 2.0 - 1.0
gt_confidence_map = (utils.to_torch(gt_confidence_map[None]) /
255.0) * 2.0 - 1.0
# exec(utils.TEST_EMBEDDING)
result_dict = {
keys.SPARSE_DEPTH: sparse_img,
keys.NN_FILLED_SPARSE_DEPTH: sparse_nn_depth,
keys.NN_FILLED_SPARSE_CONFIDENCE: sparse_confidence_map,
keys.ANNOTATED_DEPTH: gt_img,
keys.NN_FILLED_ANNOTATED_DEPTH: gt_nn_depth,
keys.NN_FILLED_ANNOTATED_CONFIDENCE: gt_confidence_map,
# 'sparse_depth_path': str(self.sparse_depth_paths[idx]),
# 'gt_path': str(self.gt_paths[idx])
}
if self.load_rgb:
rgb_img = self._read_rgb_image(str(self.rgb_paths[idx]))
result_dict[keys.ALIGNED_RGB] = rgb_img
if random.random() > self.thres:
result_dict = self.flip_result_dict(result_dict)
return result_dict
def get_homography_by_index(self, index):
homography = self.raw_data.root.homographies[index]
homography = utils.to_torch(homography)
homography = homography / homography[2:3, 2:3]
return homography