def render_cubified_voxels(voxels: torch.Tensor,
shader_type=HardPhongShader,
device="cpu",
**kwargs):
"""
Use the Cubify operator to convert inputs voxels to a mesh and then render that mesh.
Args:
voxels: FloatTensor of shape (N, D, D, D) where N is the batch size and
D is the number of voxels along each dimension.
shader_type: shader_type: shader_type: Shader to use for rendering. Examples
include HardPhongShader (default), SoftPhongShader etc or any other type
of valid Shader class.
device: torch.device on which the tensors should be located.
**kwargs: Accepts any of the kwargs that the renderer supports.
Returns:
Batch of rendered images of shape (N, H, W, 3).
"""
cubified_voxels = cubify(voxels, CUBIFY_THRESH).to(device)
cubified_voxels.textures = TexturesVertex(verts_features=torch.ones_like(
cubified_voxels.verts_padded(), device=device))
cameras = BlenderCamera(device=device)
renderer = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=kwargs.get("raster_settings",
RasterizationSettings()),
),
shader=shader_type(
device=device,
cameras=cameras,
lights=kwargs.get("lights", PointLights()).to(device),
),
)
return renderer(cubified_voxels)
def run_on_images(self, imgs, sid, mid, iid, sampled_idx):
dir1 = os.path.join(output_dir, str(sid), str(mid))
if not os.path.exists(dir1):
os.makedirs(dir1)
deprocess = imagenet_deprocess(rescale_image=False)
image_features = self.encoder(imgs)
raw_features, generated_volume = self.decoder(image_features)
generated_volume = self.merger(raw_features, generated_volume)
generated_volume = self.refiner(generated_volume)
mesh = cubify(generated_volume, 0.3)
# mesh = voxel_to_world(meshes)
save_mesh = os.path.join(dir1, "%s_%s.obj" % (iid, sampled_idx))
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(save_mesh, verts, faces)
generated_volume = generated_volume.squeeze()
img = image_to_numpy(deprocess(imgs[0][0]))
save_img = os.path.join(dir1, "%02d.png" % (iid))
# cv2.imwrite(save_img, img[:, :, ::-1])
cv2.imwrite(save_img, img)
img1 = image_to_numpy(deprocess(imgs[0][1]))
save_img1 = os.path.join(dir1, "%02d.png" % (sampled_idx))
cv2.imwrite(save_img1, img1)
# cv2.imwrite(save_img1, img1[:, :, ::-1])
get_volume_views(generated_volume, dir1, iid, sampled_idx)
def cubify(self, voxel_scores):
V = self.voxel_size
N = voxel_scores.shape[0]
voxel_probs = voxel_scores.sigmoid()
active_voxels = voxel_probs > self.cubify_threshold
voxels_per_mesh = (active_voxels.view(N, -1).sum(dim=1)).tolist()
start = V // 4
stop = start + V // 2
for i in range(N):
if voxels_per_mesh[i] == 0:
voxel_probs[i, start:stop, start:stop, start:stop] = 1
meshes = cubify(voxel_probs, self.cubify_threshold)
meshes = self._add_dummies(meshes)
meshes = voxel_to_world(meshes)
return meshes
def test_align(self):
N, V = 1, 2
device = torch.device("cuda:0")
voxels = torch.ones((N, V, V, V), dtype=torch.float32, device=device)
# topleft align
mesh = cubify(voxels, 0.5)
verts, faces = mesh.get_mesh_verts_faces(0)
self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
self.assertClose(verts.max(), torch.tensor(3.0, device=device))
# corner align
mesh = cubify(voxels, 0.5, align="corner")
verts, faces = mesh.get_mesh_verts_faces(0)
self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
self.assertClose(verts.max(), torch.tensor(1.0, device=device))
# center align
mesh = cubify(voxels, 0.5, align="center")
verts, faces = mesh.get_mesh_verts_faces(0)
self.assertClose(verts.min(), torch.tensor(-2.0, device=device))
self.assertClose(verts.max(), torch.tensor(2.0, device=device))
# invalid align
with self.assertRaisesRegex(ValueError, "Align mode must be one of"):
cubify(voxels, 0.5, align="")
# invalid align
with self.assertRaisesRegex(ValueError, "Align mode must be one of"):
cubify(voxels, 0.5, align="topright")
# inside occupancy, similar to GH#185 use case
N, V = 1, 4
voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
voxels[0, : V // 2, : V // 2, : V // 2] = 1.0
mesh = cubify(voxels, 0.5, align="corner")
verts, faces = mesh.get_mesh_verts_faces(0)
self.assertClose(verts.min(), torch.tensor(-1.0, device=device))
self.assertClose(verts.max(), torch.tensor(0.0, device=device))
def convert():
cubify(voxels, 0.5)
torch.cuda.synchronize()
def test_cubify(self):
N, V = 4, 2
device = torch.device("cuda:0")
voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
# 1st example: (top left corner, znear) is on
voxels[0, 0, 0, 0] = 1.0
# 2nd example: all are on
voxels[1] = 1.0
# 3rd example: empty
# 4th example
voxels[3, :, :, 1] = 1.0
voxels[3, 1, 1, 0] = 1.0
# compute cubify
meshes = cubify(voxels, 0.5)
# 1st-check
verts, faces = meshes.get_mesh_verts_faces(0)
self.assertClose(faces.max().cpu(), torch.tensor(verts.size(0) - 1))
self.assertClose(
verts,
torch.tensor(
[
[-1.0, -1.0, -1.0],
[-1.0, -1.0, 1.0],
[1.0, -1.0, -1.0],
[1.0, -1.0, 1.0],
[-1.0, 1.0, -1.0],
[-1.0, 1.0, 1.0],
[1.0, 1.0, -1.0],
[1.0, 1.0, 1.0],
],
dtype=torch.float32,
device=device,
),
)
self.assertClose(
faces,
torch.tensor(
[
[0, 1, 4],
[1, 5, 4],
[4, 5, 6],
[5, 7, 6],
[0, 4, 6],
[0, 6, 2],
[0, 3, 1],
[0, 2, 3],
[6, 7, 3],
[6, 3, 2],
[1, 7, 5],
[1, 3, 7],
],
dtype=torch.int64,
device=device,
),
)
# 2nd-check
verts, faces = meshes.get_mesh_verts_faces(1)
self.assertClose(faces.max().cpu(), torch.tensor(verts.size(0) - 1))
self.assertClose(
verts,
torch.tensor(
[
[-1.0, -1.0, -1.0],
[-1.0, -1.0, 1.0],
[-1.0, -1.0, 3.0],
[1.0, -1.0, -1.0],
[1.0, -1.0, 1.0],
[1.0, -1.0, 3.0],
[3.0, -1.0, -1.0],
[3.0, -1.0, 1.0],
[3.0, -1.0, 3.0],
[-1.0, 1.0, -1.0],
[-1.0, 1.0, 1.0],
[-1.0, 1.0, 3.0],
[1.0, 1.0, -1.0],
[1.0, 1.0, 3.0],
[3.0, 1.0, -1.0],
[3.0, 1.0, 1.0],
[3.0, 1.0, 3.0],
[-1.0, 3.0, -1.0],
[-1.0, 3.0, 1.0],
[-1.0, 3.0, 3.0],
[1.0, 3.0, -1.0],
[1.0, 3.0, 1.0],
[1.0, 3.0, 3.0],
[3.0, 3.0, -1.0],
[3.0, 3.0, 1.0],
[3.0, 3.0, 3.0],
],
dtype=torch.float32,
device=device,
),
)
self.assertClose(
faces,
torch.tensor(
[
[0, 1, 9],
[1, 10, 9],
[0, 9, 12],
[0, 12, 3],
[0, 4, 1],
[0, 3, 4],
[1, 2, 10],
[2, 11, 10],
[1, 5, 2],
[1, 4, 5],
[2, 13, 11],
[2, 5, 13],
[3, 12, 14],
[3, 14, 6],
[3, 7, 4],
[3, 6, 7],
[14, 15, 7],
[14, 7, 6],
[4, 8, 5],
[4, 7, 8],
[15, 16, 8],
[15, 8, 7],
[5, 16, 13],
[5, 8, 16],
[9, 10, 17],
[10, 18, 17],
[17, 18, 20],
[18, 21, 20],
[9, 17, 20],
[9, 20, 12],
[10, 11, 18],
[11, 19, 18],
[18, 19, 21],
[19, 22, 21],
[11, 22, 19],
[11, 13, 22],
[20, 21, 23],
[21, 24, 23],
[12, 20, 23],
[12, 23, 14],
[23, 24, 15],
[23, 15, 14],
[21, 22, 24],
[22, 25, 24],
[24, 25, 16],
[24, 16, 15],
[13, 25, 22],
[13, 16, 25],
],
dtype=torch.int64,
device=device,
),
)
# 3rd-check
verts, faces = meshes.get_mesh_verts_faces(2)
self.assertTrue(verts.size(0) == 0)
self.assertTrue(faces.size(0) == 0)
# 4th-check
verts, faces = meshes.get_mesh_verts_faces(3)
self.assertClose(
verts,
torch.tensor(
[
[1.0, -1.0, -1.0],
[1.0, -1.0, 1.0],
[1.0, -1.0, 3.0],
[3.0, -1.0, -1.0],
[3.0, -1.0, 1.0],
[3.0, -1.0, 3.0],
[-1.0, 1.0, 1.0],
[-1.0, 1.0, 3.0],
[1.0, 1.0, -1.0],
[1.0, 1.0, 1.0],
[1.0, 1.0, 3.0],
[3.0, 1.0, -1.0],
[3.0, 1.0, 1.0],
[3.0, 1.0, 3.0],
[-1.0, 3.0, 1.0],
[-1.0, 3.0, 3.0],
[1.0, 3.0, -1.0],
[1.0, 3.0, 1.0],
[1.0, 3.0, 3.0],
[3.0, 3.0, -1.0],
[3.0, 3.0, 1.0],
[3.0, 3.0, 3.0],
],
dtype=torch.float32,
device=device,
),
)
self.assertClose(
faces,
torch.tensor(
[
[0, 1, 8],
[1, 9, 8],
[0, 8, 11],
[0, 11, 3],
[0, 4, 1],
[0, 3, 4],
[11, 12, 4],
[11, 4, 3],
[1, 2, 9],
[2, 10, 9],
[1, 5, 2],
[1, 4, 5],
[12, 13, 5],
[12, 5, 4],
[2, 13, 10],
[2, 5, 13],
[6, 7, 14],
[7, 15, 14],
[14, 15, 17],
[15, 18, 17],
[6, 14, 17],
[6, 17, 9],
[6, 10, 7],
[6, 9, 10],
[7, 18, 15],
[7, 10, 18],
[8, 9, 16],
[9, 17, 16],
[16, 17, 19],
[17, 20, 19],
[8, 16, 19],
[8, 19, 11],
[19, 20, 12],
[19, 12, 11],
[17, 18, 20],
[18, 21, 20],
[20, 21, 13],
[20, 13, 12],
[10, 21, 18],
[10, 13, 21],
],
dtype=torch.int64,
device=device,
),
)
def test_allempty(self):
N, V = 32, 14
device = torch.device("cuda:0")
voxels = torch.zeros((N, V, V, V), dtype=torch.float32, device=device)
meshes = cubify(voxels, 0.5)
self.assertTrue(meshes.isempty())
def _forward_shape(self, features, instances):
"""
Forward logic for the voxel and mesh refinement branch.
Args:
features (list[Tensor]): #level input features for voxel prediction
instances (list[Instances]): the per-image instances to train/predict meshes.
In training, they can be the proposals.
In inference, they can be the predicted boxes.
Returns:
In training, a dict of losses.
In inference, update `instances` with new fields "pred_voxels" & "pred_meshes" and return it.
"""
if not self.voxel_on and not self.mesh_on:
return {} if self.training else instances
if self.training:
# The loss is only defined on positive proposals.
proposals, _ = select_foreground_proposals(instances,
self.num_classes)
proposal_boxes = [x.proposal_boxes for x in proposals]
losses = {}
if self.voxel_on:
voxel_features = self.voxel_pooler(features, proposal_boxes)
voxel_logits = self.voxel_head(voxel_features)
loss_voxel, target_voxels = voxel_rcnn_loss(
voxel_logits,
proposals,
loss_weight=self.voxel_loss_weight)
losses.update({"loss_voxel": loss_voxel})
if self._vis:
self._misc["target_voxels"] = target_voxels
if self.cls_agnostic_voxel:
with torch.no_grad():
vox_in = voxel_logits.sigmoid().squeeze(
1) # (N, V, V, V)
init_mesh = cubify(vox_in, self.cubify_thresh) # 1
else:
raise ValueError(
"No support for class specific predictions")
if self.mesh_on:
mesh_features = self.mesh_pooler(features, proposal_boxes)
if not self.voxel_on:
if mesh_features.shape[0] > 0:
init_mesh = ico_sphere(self.ico_sphere_level,
mesh_features.device)
init_mesh = init_mesh.extend(mesh_features.shape[0])
else:
init_mesh = Meshes(verts=[], faces=[])
pred_meshes = self.mesh_head(mesh_features, init_mesh)
# loss weights
loss_weights = {
"chamfer": self.chamfer_loss_weight,
"normals": self.normals_loss_weight,
"edge": self.edge_loss_weight,
}
if not pred_meshes[0].isempty():
loss_chamfer, loss_normals, loss_edge, target_meshes = mesh_rcnn_loss(
pred_meshes,
proposals,
loss_weights=loss_weights,
gt_num_samples=self.gt_num_samples,
pred_num_samples=self.pred_num_samples,
gt_coord_thresh=self.gt_coord_thresh,
)
if self._vis:
self._misc["init_meshes"] = init_mesh
self._misc["target_meshes"] = target_meshes
else:
loss_chamfer = sum(
k.sum() for k in self.mesh_head.parameters()) * 0.0
loss_normals = sum(
k.sum() for k in self.mesh_head.parameters()) * 0.0
loss_edge = sum(k.sum()
for k in self.mesh_head.parameters()) * 0.0
losses.update({
"loss_chamfer": loss_chamfer,
"loss_normals": loss_normals,
"loss_edge": loss_edge,
})
return losses
else:
pred_boxes = [x.pred_boxes for x in instances]
if self.voxel_on:
voxel_features = self.voxel_pooler(features, pred_boxes)
voxel_logits = self.voxel_head(voxel_features)
voxel_rcnn_inference(voxel_logits, instances)
if self.cls_agnostic_voxel:
with torch.no_grad():
vox_in = voxel_logits.sigmoid().squeeze(
1) # (N, V, V, V)
init_mesh = cubify(vox_in, self.cubify_thresh) # 1
else:
raise ValueError(
"No support for class specific predictions")
if self.mesh_on:
mesh_features = self.mesh_pooler(features, pred_boxes)
if not self.voxel_on:
if mesh_features.shape[0] > 0:
init_mesh = ico_sphere(self.ico_sphere_level,
mesh_features.device)
init_mesh = init_mesh.extend(mesh_features.shape[0])
else:
init_mesh = Meshes(verts=[], faces=[])
pred_meshes = self.mesh_head(mesh_features, init_mesh)
mesh_rcnn_inference(pred_meshes[-1], instances)
else:
assert self.voxel_on
mesh_rcnn_inference(init_mesh, instances)
return instances
