def handle_model()

in tools/preprocess_shapenet.py [0:0]

• 66 lines of code
• 11 McCabe index (conditional complexity)

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
    # Some ShapeNet textures are corrupt. Since we don't need them,
    # avoid loading textures altogether.
    mesh = load_obj(obj_path, load_textures=False)
    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)