def get_camera(self, opt, idx):
c, m, i, v = self.list[idx]
cam_fname = "{0}/{1}/{2}/cameras.npz".format(self.path, c, m)
cam = np.load(cam_fname)
focal = 1.8660254037844388
intr = torch.tensor([[focal * opt.W, 0, opt.W / 2],
[0, focal * opt.H, opt.H / 2], [0, 0, 1]])
extr = torch.from_numpy(cam["world_mat_{}".format(v)][:3]).float()
pose = camera.pose(R=extr[:, :3], t=extr[:, 3])
return intr, pose
def __getitem__(self, idx):
opt = self.opt
name, mask_id = self.list[idx]
if name == "n02690373_10374": # no mask data
return self[np.random.randint(len(self))]
sample = dict(idx=idx)
aug = self.generate_augmentation(opt) if self.augment else None
# load camera
meta = self.get_metadata(opt, idx)
pose_cam = camera.pose(t=[0, 0, opt.camera.dist])
assert (opt.camera.model == "orthographic")
pose = self.cameras[idx] if opt.data.preload else self.get_camera(
opt, idx, meta=meta)
pose = camera.pose.compose([pose, pose_cam])
if aug is not None:
pose = self.augment_camera(opt, pose, aug, pose_cam=pose_cam)
intr = False # there are no None tensors
sample.update(
pose=pose,
intr=intr,
)
# load images and compute distance transform
image = self.images[idx] if opt.data.preload else self.get_image(
opt, idx)
rgb, mask, crop_size = self.preprocess_image(opt,
image,
bbox=meta.bbox,
aug=aug)
meta.crop_size = crop_size
dt = self.compute_dist_transform(opt, mask)
sample.update(
rgb_input_map=rgb,
mask_input_map=mask,
dt_input_map=dt,
)
# vectorize images (and randomly sample)
rgb = rgb.permute(1, 2, 0).view(opt.H * opt.W, 3)
mask = mask.permute(1, 2, 0).view(opt.H * opt.W, 1)
dt = dt.permute(1, 2, 0).view(opt.H * opt.W, 1)
if self.split == "train" and opt.impl.rand_sample:
ray_idx = torch.randperm(opt.H * opt.W)[:opt.impl.rand_sample]
rgb, mask, dt = rgb[ray_idx], mask[ray_idx], dt[ray_idx]
sample.update(ray_idx=ray_idx)
sample.update(
rgb_input=rgb,
mask_input=mask,
dt_input=dt,
)
# load GT point cloud (only for validation!)
dpc = self.pointclouds[
idx] if opt.data.preload else self.get_pointcloud(
opt, idx, meta=meta)
sample.update(dpc=dpc)
return sample
def augment_camera(self, opt, pose, aug, pose_cam=None):
if aug.flip:
raise NotImplementedError
if aug.rot_angle:
angle = torch.tensor(aug.rot_angle) * np.pi / 180
R = camera.angle_to_rotation_matrix(-angle,
axis="Z") # in-plane rotation
rot_inplane = camera.pose(R=R)
pose = camera.pose.compose(
[pose,
camera.pose.invert(pose_cam), rot_inplane, pose_cam])
return pose
def get_camera(self, opt, idx, meta=None):
azim = torch.tensor(meta.cam.azim / 180 * np.pi).float()
elev = torch.tensor(meta.cam.elev / 180 * np.pi).float()
theta = torch.tensor(meta.cam.theta / 180 * np.pi).float()
Rz = camera.angle_to_rotation_matrix(azim, axis="Z")
Rx = camera.angle_to_rotation_matrix(elev - np.pi / 2, axis="X")
R2d = camera.angle_to_rotation_matrix(-theta, axis="Z")
R = R2d @ Rx @ Rz
R_trans = torch.tensor([[0, 0, 1], [-1, 0, 0], [0, 1, 0]],
dtype=torch.float32)
pose = camera.pose(R=R @ R_trans)
return pose
def __getitem__(self, idx):
opt = self.opt
cad_idx = self.list[idx][2]
sample = dict(
idx=idx,
cad_idx=cad_idx,
)
aug = self.generate_augmentation(opt) if self.augment else None
# load camera
pose_cam = camera.pose(R=[[-1, 0, 0], [0, -1, 0], [0, 0, -1]],
t=[0, 0, opt.camera.dist])
intr, pose = self.cameras[
idx] if opt.data.preload else self.get_camera(opt, idx)
if aug is not None:
pose = self.augment_camera(opt, pose, aug, pose_cam=pose_cam)
sample.update(
pose=pose,
intr=intr,
)
# load images and compute distance transform
image = self.images[idx] if opt.data.preload else self.get_image(
opt, idx)
rgb, mask = self.preprocess_image(opt, image, aug=aug)
dt = self.compute_dist_transform(opt, mask, intr=intr)
sample.update(
rgb_input_map=rgb,
mask_input_map=mask,
dt_input_map=dt,
)
# vectorize images (and randomly sample)
rgb = rgb.permute(1, 2, 0).view(opt.H * opt.W, 3)
mask = mask.permute(1, 2, 0).view(opt.H * opt.W, 1)
dt = dt.permute(1, 2, 0).view(opt.H * opt.W, 1)
if self.split == "train" and opt.impl.rand_sample:
ray_idx = torch.randperm(opt.H * opt.W)[:opt.impl.rand_sample]
rgb, mask, dt = rgb[ray_idx], mask[ray_idx], dt[ray_idx]
sample.update(ray_idx=ray_idx)
sample.update(
rgb_input=rgb,
mask_input=mask,
dt_input=dt,
)
# load GT point cloud (only for validation!)
dpc = self.pointclouds[
idx] if opt.data.preload else self.get_pointcloud(opt, idx)
sample.update(dpc=dpc)
return sample