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