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 save(
self,
data: Pointclouds,
path: Union[str, Path],
path_manager: PathManager,
binary: Optional[bool],
decimal_places: Optional[int] = None,
**kwargs,
) -> bool:
if not endswith(path, self.known_suffixes):
return False
points = data.points_list()[0]
features = data.features_packed()
normals = data.normals_packed()
with _open_file(path, path_manager, "wb") as f:
_save_ply(
f=f,
verts=points,
verts_colors=features,
verts_normals=normals,
faces=None,
ascii=binary is False,
decimal_places=decimal_places,
)
return True
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 _add_pointcloud_trace(
fig: go.Figure,
pointclouds: Pointclouds,
trace_name: str,
subplot_idx: int,
ncols: int,
max_points_per_pointcloud: int,
marker_size: int,
): # pragma: no cover
"""
Adds a trace rendering a Pointclouds object to the passed in figure, with
a given name and in a specific subplot.
Args:
fig: plotly figure to add the trace within.
pointclouds: Pointclouds object to render. It can be batched.
trace_name: name to label the trace with.
subplot_idx: identifies the subplot, with 0 being the top left.
ncols: the number of subplots per row.
max_points_per_pointcloud: the number of points to render, which are randomly sampled.
marker_size: the size of the rendered points
"""
pointclouds = pointclouds.detach().cpu().subsample(
max_points_per_pointcloud)
verts = pointclouds.points_packed()
features = pointclouds.features_packed()
color = None
if features is not None:
if features.shape[1] == 4: # rgba
template = "rgb(%d, %d, %d, %f)"
rgb = (features[:, :3].clamp(0.0, 1.0) * 255).int()
color = [
template % (*rgb_, a_)
for rgb_, a_ in zip(rgb, features[:, 3])
]
if features.shape[1] == 3:
template = "rgb(%d, %d, %d)"
rgb = (features.clamp(0.0, 1.0) * 255).int()
color = [template % (r, g, b) for r, g, b in rgb]
row = subplot_idx // ncols + 1
col = subplot_idx % ncols + 1
fig.add_trace(
go.Scatter3d(
x=verts[:, 0],
y=verts[:, 1],
z=verts[:, 2],
marker={
"color": color,
"size": marker_size
},
mode="markers",
name=trace_name,
),
row=row,
col=col,
)
# Access the current subplot's scene configuration
plot_scene = "scene" + str(subplot_idx + 1)
current_layout = fig["layout"][plot_scene]
# update the bounds of the axes for the current trace
verts_center = verts.mean(0)
max_expand = (verts.max(0)[0] - verts.min(0)[0]).max()
_update_axes_bounds(verts_center, max_expand, current_layout)
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 _add_pointcloud_trace(
fig: go.Figure,
pointclouds: Pointclouds,
trace_name: str,
subplot_idx: int,
ncols: int,
max_points_per_pointcloud: int,
marker_size: int,
):
"""
Adds a trace rendering a Pointclouds object to the passed in figure, with
a given name and in a specific subplot.
Args:
fig: plotly figure to add the trace within.
pointclouds: Pointclouds object to render. It can be batched.
trace_name: name to label the trace with.
subplot_idx: identifies the subplot, with 0 being the top left.
ncols: the number of sublpots per row.
max_points_per_pointcloud: the number of points to render, which are randomly sampled.
marker_size: the size of the rendered points
"""
pointclouds = pointclouds.detach().cpu()
verts = pointclouds.points_packed()
features = pointclouds.features_packed()
indices = None
if pointclouds.num_points_per_cloud().max() > max_points_per_pointcloud:
start_index = 0
index_list = []
for num_points in pointclouds.num_points_per_cloud():
if num_points > max_points_per_pointcloud:
indices_cloud = np.random.choice(num_points,
max_points_per_pointcloud,
replace=False)
index_list.append(start_index + indices_cloud)
else:
index_list.append(start_index + np.arange(num_points))
start_index += num_points
indices = np.concatenate(index_list)
verts = verts[indices]
color = None
if features is not None:
if indices is not None:
# Only select features if we selected vertices above
features = features[indices]
if features.shape[1] == 4: # rgba
template = "rgb(%d, %d, %d, %f)"
rgb = (features[:, :3].clamp(0.0, 1.0) * 255).int()
color = [
template % (*rgb_, a_)
for rgb_, a_ in zip(rgb, features[:, 3])
]
if features.shape[1] == 3:
template = "rgb(%d, %d, %d)"
rgb = (features.clamp(0.0, 1.0) * 255).int()
color = [template % (r, g, b) for r, g, b in rgb]
row = subplot_idx // ncols + 1
col = subplot_idx % ncols + 1
fig.add_trace(
go.Scatter3d( # pyre-ignore[16]
x=verts[:, 0],
y=verts[:, 1],
z=verts[:, 2],
marker={
"color": color,
"size": marker_size
},
mode="markers",
name=trace_name,
),
row=row,
col=col,
)
# Access the current subplot's scene configuration
plot_scene = "scene" + str(subplot_idx + 1)
current_layout = fig["layout"][plot_scene]
# update the bounds of the axes for the current trace
verts_center = verts.mean(0)
max_expand = (verts.max(0)[0] - verts.min(0)[0]).max()
_update_axes_bounds(verts_center, max_expand, current_layout)
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))