def forward(self, pts, features):
# Make sure pts and features are equal
assert pts.size(2) == 3
assert pts.size(1) == features.size(1)
pts[:, :, 0] = -pts[:, :, 0]
pts[:, :, 1] = -pts[:, :, 1]
radius = float(self.radius) / float(self.img_size) * 2.0
params = compositing.CompositeParams(radius=radius)
pointcloud = Pointclouds(points=pts, features=features)
points_idx, _, dist = rasterize_points(pointcloud, self.img_size,
radius, self.points_per_pixel)
dist = dist / pow(radius, 2)
alphas = (1 - dist.clamp(max=1, min=1e-3).pow(0.5).pow(
self.gamma).permute(0, 3, 1, 2))
transformed_feature_alphas = compositing.alpha_composite(
points_idx.permute(0, 3, 1, 2).long(), alphas,
pointcloud.features_packed().permute(1, 0), params)
return transformed_feature_alphas
def forward(self, pts3D, src):
bs = src.size(0)
if len(src.size()) > 3:
bs, c, w, _ = src.size()
image_size = w
pts3D = pts3D.permute(0, 2, 1)
src = src.unsqueeze(2).repeat(1, 1, w, 1, 1).view(bs, c, -1)
else:
bs = src.size(0)
image_size = self.size
# Make sure these have been arranged in the same way
assert pts3D.size(2) == 3
assert pts3D.size(1) == src.size(2)
pts3D[:, :, 1] = -pts3D[:, :, 1]
pts3D[:, :, 0] = -pts3D[:, :, 0]
# Add on the default feature to the end of the src
# src = torch.cat((src, self.default_feature.repeat(bs, 1, 1)), 2)
radius = float(self.radius) / float(image_size) * 2.0
pts3D = Pointclouds(points=pts3D, features=src.permute(0, 2, 1))
points_idx, _, dist = rasterize_points(pts3D, image_size, radius,
self.points_per_pixel)
if os.environ["DEBUG"]:
print("Max dist: ", dist.max(), pow(radius, self.opts.rad_pow))
dist = dist / pow(radius, self.opts.rad_pow)
if os.environ["DEBUG"]:
print("Max dist: ", dist.max())
alphas = ((1 - dist.clamp(max=1, min=1e-3).pow(0.5)).pow(
self.opts.tau).permute(0, 3, 1, 2))
if self.opts.accumulation == 'alphacomposite':
transformed_src_alphas = compositing.alpha_composite(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
elif self.opts.accumulation == 'wsum':
transformed_src_alphas = compositing.weighted_sum(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
elif self.opts.accumulation == 'wsumnorm':
transformed_src_alphas = compositing.weighted_sum_norm(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
return transformed_src_alphas
def forward(self, points, features) -> torch.Tensor:
"""
points FloatTensor B x N x 3
features FloatTensor B x N x F
"""
# Rasterize points -- bins set heurisically
pointcloud = pytorch3d.structures.Pointclouds(points, features=features)
idx, zbuf, dist_xy = rasterize_points(pointcloud, self.S, self.r, self.K)
# Calculate PC coverage
valid_pts = (idx >= 0).float()
valid_ray = valid_pts[:, :, :, 0]
# Calculate composite weights -- dist_xy is squared distance!!
# Clamp weights to avoid 0 gradients or errors
weights = self.calculate_weights(dist_xy, self.r)
weights = weights.clamp(min=0.0, max=0.99)
# Composite the raster for feats and depth
idx = idx.long().permute(0, 3, 1, 2).contiguous()
feats = pointcloud.features_packed().permute(1, 0)
feats = self.compositor(idx, weights, feats)
# == Rasterize depth ==
# zero out weights -- currently applies norm_weighted sum
w_normed = weights * (idx >= 0).float()
w_normed = w_normed / w_normed.sum(dim=1, keepdim=True).clamp(min=1e-9)
z_weighted = zbuf.permute(0, 3, 1, 2).contiguous() * w_normed.contiguous()
z_weighted = z_weighted.sum(dim=1, keepdim=True)
return {
"raster_output": {
"idx": idx,
"zbuf": zbuf,
"dist_xy": dist_xy,
"alphas": weights,
"points": points,
"feats": features,
},
"feats": feats,
"depth": z_weighted,
"mask": valid_ray,
"valid_rays": valid_ray.mean(dim=(1, 2)),
"valid_pts": valid_pts.mean(dim=(1, 2, 3)),
}
def forward(self, pts3D, src):
bs = src.size(0)
if len(src.size()) > 3:
bs, c, w, _ = src.size()
image_size = w
pts3D = pts3D.permute(0, 2, 1)
src = src.unsqueeze(2).repeat(1, 1, w, 1, 1).view(bs, c, -1)
else:
bs = src.size(0)
image_size = self.size
# Make sure these have been arranged in the same way
assert pts3D.size(2) == 3
assert pts3D.size(1) == src.size(2)
# pts3D.shape = (bs, w*h, 3) --> (x,y,z) coordinate for ever element in the image raster
# Because we have done re-projection, the i-th coordinate in the image raster must no longer be identical to (x,y)!
# src.shape = (bs, c, w*h) --> c features for every element in the image raster (w*h)
#print("Features: {}".format(src.shape))
#print("3D Pointcloud: {}".format(pts3D.shape))
# flips the x and y coordinate
pts3D[:, :, 1] = -pts3D[:, :, 1]
pts3D[:, :, 0] = -pts3D[:, :, 0]
# Add on the default feature to the end of the src
#src = torch.cat((src, self.default_feature.repeat(bs, 1, 1)), 2)
radius = float(self.radius) / float(
image_size
) * 2.0 # convert radius to fit the [-1,1] NDC ?? Or is this just arbitrary scaling s.t. radius as meaningful size?
params = compositing.CompositeParams(radius=radius)
#print("Radius - before: {}, converted: {}".format(self.radius, radius))
pts3D = Pointclouds(points=pts3D, features=src.permute(0, 2, 1))
points_idx, _, dist = rasterize_points(
pts3D, image_size, radius, self.points_per_pixel
) # see method signature for meaning of these output values
#print("points_idx: {}".format(points_idx.shape))
#print("dist: {}".format(points_idx.shape))
#print("Max dist: ", dist.max(), pow(radius, self.rad_pow))
dist = dist / pow(
radius, self.rad_pow
) # equation 1 from the paper (3.2): this calculates N(p_i, l_mn) from the d2 dist
#print("Max dist: ", dist.max())
alphas = (
(1 - dist.clamp(max=1, min=1e-3).pow(0.5)).pow(
self.accumulation_tau).permute(0, 3, 1, 2)
) # equation 2 from the paper (3.2): prepares alpha values for composition of the feature vectors
#print("alphas: ", alphas.shape)
#print("pointclouds object: {}".format(pts3D.features_packed().shape))
#print("alphas: ", alphas)
if self.accumulation == 'alphacomposite':
transformed_src_alphas = compositing.alpha_composite(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(
1, 0
), # pts3D also contains features here, because this is now a Pointclouds object
params,
)
elif self.accumulation == 'wsum':
transformed_src_alphas = compositing.weighted_sum(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
params,
)
elif self.accumulation == 'wsumnorm':
transformed_src_alphas = compositing.weighted_sum_norm(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
params,
)
else:
raise NotImplementedError('Unsupported accumulation type: ' +
self.accumulation)
return transformed_src_alphas
def generate_video(im_name, normalise):
# Load in the correspondences and images
im1 = Image.open(os.environ['BASE_PATH'] + '/imL/%s.jpg' % im_name)
im2 = Image.open(os.environ['BASE_PATH'] + '/imR/%s.jpg' % im_name)
if normalise:
np_tempwarp = np.load(os.environ['BASE_PATH'] +
'/warps/temp_sampler%s_2_grad_norm.npz' %
im_name)
H = np_tempwarp['H']
np_tempwarp = np_tempwarp['sampler']
else:
np_tempwarp = np.load(os.environ['BASE_PATH'] +
'/warps/temp_sampler%s_2_grad_coarse.npz.npy' %
im_name)
if normalise:
im1_arr = np.array(im1)
im2_arr = np.array(im2)
K1 = np.eye(3)
K1[0, 0] = 2 / im1_arr.shape[1]
K1[1, 1] = 2 / im1_arr.shape[0]
K1[0:2, 2] = -1
K2 = np.eye(3)
K2[0, 0] = 2 / im2_arr.shape[1]
K2[1, 1] = 2 / im2_arr.shape[0]
K2[0:2, 2] = -1
aff_mat = (np.linalg.inv(K2) @ H @ K1)
# Now transform the image and return
warp_im1 = cv2.warpAffine(im1_arr, aff_mat[0:2],
(im2_arr.shape[1], im2_arr.shape[0]))
im1 = warp_im1
im1 = Image.fromarray(im1)
warp = torch.Tensor(np_tempwarp).unsqueeze(0)
im1_torch = tr.ToTensor()(im1).unsqueeze(0)
im2_torch = tr.ToTensor()(im2).unsqueeze(0)
gen_img = F.grid_sample(im1_torch, warp)
sampler = F.upsample(warp.permute(0, 3, 1, 2),
size=(im2_torch.size(2), im2_torch.size(3)))
gen_imglarge = F.grid_sample(im1_torch, sampler.permute(0, 2, 3, 1))
W1, W2, _ = np_tempwarp.shape
orig_warp = torch.meshgrid(torch.linspace(-1, 1, W1),
torch.linspace(-1, 1, W2))
orig_warp = torch.cat(
(orig_warp[1].unsqueeze(2), orig_warp[0].unsqueeze(2)), 2)
orig_warp = orig_warp.unsqueeze(0)
warp = torch.Tensor(np_tempwarp).unsqueeze(0)
new_imgs = []
if not os.path.exists('./temp%s/%s' % (im_name, im_name)):
os.makedirs('./temp%s/%s' % (im_name, im_name))
radius = 2 * 2 / 1024.
for i in tqdm(range(-10, 30)):
resample = (orig_warp * float(i) / 20. + warp * float(20 - i) / 20.)
pts3D = resample.view(1, -1, 2)
pts_mask = (warp.view(-1, 2)[:, 0] > -1) & (warp.view(-1, 2)[:, 0] < 1)
pts3D = pts3D[:, pts_mask, :]
pts3D = -pts3D
pts3D = torch.cat((pts3D.cuda(), torch.ones(
(1, pts3D.size(1), 1)).cuda()), 2)
rgb = F.grid_sample(im2_torch,
orig_warp).permute(0, 2, 3, 1).view(1, -1,
3)[:,
pts_mask, :]
mask = torch.ones((1, rgb.size(1), 1)).cuda()
pts3DRGB = Pointclouds(points=pts3D, features=rgb)
points_idx, _, dist = rasterize_points(pts3DRGB, 1024, radius, 1)
gen_img = pts3DRGB.features_packed()[
points_idx.permute(0, 3, 1, 2).long()[0], :].permute(
0, 3, 1, 2).mean(dim=0, keepdim=True)
new_imgs += [gen_img.squeeze().permute(1, 2, 0)]
torchvision.utils.save_image(
gen_img, './temp%s/%s/im-%03d.png' % (im_name, im_name, i + 10))
mask = (points_idx.permute(0, 3, 1, 2) < 0).float()
torchvision.utils.save_image(
mask, './temp%s/%s/mask-%03d.png' % (im_name, im_name, i + 10))