def setup(self, cfg):
# fmt: off
self.cubify_threshold = cfg.MODEL.VOXEL_HEAD.CUBIFY_THRESH
self.voxel_size = cfg.MODEL.VOXEL_HEAD.VOXEL_SIZE
# fmt: on
self.register_buffer("K", get_blender_intrinsic_matrix())
def __init__(self, cfg):
nn.Module.__init__(self)
self.ico_sphere_level = cfg.MODEL.MESH_HEAD.ICO_SPHERE_LEVEL
self.init_mvsnet(cfg)
self.init_mesh_head(cfg)
self.cfg = cfg
self.register_buffer("K", get_blender_intrinsic_matrix())
def __init__(self, cfg):
super(Pixel2MeshHead, self).__init__()
# fmt: off
backbone = cfg.MODEL.MESH_HEAD.RGB_BACKBONE
self.ico_sphere_level = cfg.MODEL.MESH_HEAD.ICO_SPHERE_LEVEL
# fmt: on
self.register_buffer("K", get_blender_intrinsic_matrix())
# backbone
self.backbone, feat_dims = build_backbone(backbone)
# mesh head
cfg.MODEL.MESH_HEAD.COMPUTED_INPUT_CHANNELS = sum(feat_dims)
self.mesh_head = MeshRefinementHead(cfg)
def __init__(self, cfg):
super(VoxMeshHead, self).__init__()
# fmt: off
backbone = cfg.MODEL.BACKBONE
self.cubify_threshold = cfg.MODEL.VOXEL_HEAD.CUBIFY_THRESH
self.voxel_size = cfg.MODEL.VOXEL_HEAD.VOXEL_SIZE
# fmt: on
self.register_buffer("K", get_blender_intrinsic_matrix())
# backbone
self.backbone, feat_dims = build_backbone(backbone)
# voxel head
cfg.MODEL.VOXEL_HEAD.COMPUTED_INPUT_CHANNELS = feat_dims[-1]
self.voxel_head = VoxelHead(cfg)
# mesh head
cfg.MODEL.MESH_HEAD.COMPUTED_INPUT_CHANNELS = sum(feat_dims)
self.mesh_head = MeshRefinementHead(cfg)
def handle_model(args, sid, mid, i, N):
if (i + 1) % args.print_every == 0:
logger.info(" Handling model %d / %d" % (i + 1, N))
shapenet_dir = os.path.join(args.shapenet_dir, sid, mid)
binvox_dir = os.path.join(args.shapenet_binvox_dir, sid, mid)
r2n2_dir = os.path.join(args.r2n2_dir, sid, mid, "rendering")
output_dir = os.path.join(args.output_dir, sid, mid)
obj_path = os.path.join(shapenet_dir, "model.obj")
if not os.path.isfile(obj_path):
logger.info("WARNING: Skipping %s/%s, no .obj" % (sid, mid))
return None
mesh = load_obj(obj_path)
verts = mesh[0]
faces = mesh[1].verts_idx
# WARNING: Here we hardcode the assumption that the input voxels are
# 128 x 128 x 128.
voxel_path = os.path.join(binvox_dir, "model.presolid.binvox")
if not os.path.isfile(voxel_path):
logger.info("WARNING: Skipping %s/%s, no voxels" % (sid, mid))
return None
with open(voxel_path, "rb") as f:
# Read voxel coordinates as a tensor of shape (N, 3)
voxel_coords = read_binvox_coords(f)
if not os.path.isdir(r2n2_dir):
logger.info("WARNING: Skipping %s/%s, no images" % (sid, mid))
return None
os.makedirs(output_dir)
# Save metadata.pt
image_list = load_image_list(args, sid, mid)
extrinsics = load_extrinsics(args, sid, mid)
intrinsic = get_blender_intrinsic_matrix()
metadata = {"image_list": image_list, "intrinsic": intrinsic, "extrinsics": extrinsics}
metadata_path = os.path.join(output_dir, "metadata.pt")
torch.save(metadata, metadata_path)
# Save mesh.pt
mesh_data = {"verts": verts, "faces": faces}
mesh_path = os.path.join(output_dir, "mesh.pt")
torch.save(mesh_data, mesh_path)
# Align voxels to the same coordinate system as mesh verts, and save voxels.pt
voxel_coords = align_bbox(voxel_coords, verts)
voxel_data = {"voxel_coords": voxel_coords}
voxel_path = os.path.join(output_dir, "voxels.pt")
torch.save(voxel_data, voxel_path)
for V in args.voxel_sizes:
voxel_dir = os.path.join(output_dir, "vox%d" % V)
if not os.path.isdir(voxel_dir):
os.makedirs(voxel_dir)
for i in range(len(image_list)):
P = intrinsic.mm(extrinsics[i]).cuda()
voxels = voxelize(voxel_coords.cuda(), P, V).cpu()
voxel_path = os.path.join(voxel_dir, "%03d.pt" % i)
torch.save(voxels, voxel_path)
if args.num_samples > 0:
mesh = Meshes(verts=[verts.cuda()], faces=[faces.cuda()])
points, normals = sample_points_from_meshes(
mesh, num_samples=args.num_samples, return_normals=True
)
points_sampled = points[0].cpu().detach()
normals_sampled = normals[0].cpu().detach()
samples_data = {"points_sampled": points_sampled, "normals_sampled": normals_sampled}
samples_path = os.path.join(output_dir, "samples.pt")
torch.save(samples_data, samples_path)
# Copy the images to the output directory
output_img_dir = os.path.join(output_dir, "images")
os.makedirs(output_img_dir)
for fn in image_list:
ext = os.path.splitext(fn)[1]
assert ext in IMAGE_EXTS
src = os.path.join(r2n2_dir, fn)
dst = os.path.join(output_img_dir, fn)
shutil.copy(src, dst)
num_imgs = len(image_list)
return (sid, mid, num_imgs)