facebook/actionbench

1"""Project animated point cloud onto image using camera parameters."""2 3import json4 5import numpy as np6from PIL import Image7from scipy.ndimage import gaussian_filter8 9 10def project_points(11    points: np.ndarray, camera: dict, img_size: int12) -> tuple[np.ndarray, np.ndarray]:13    """Project 3D points to pixel coordinates.14 15    Uses Blender camera convention:16      X_cam = points @ R^T + T, camera looks along -Z.17 18    Args:19        points: (N, 3) world-space points.20        camera: dict with R, T, focal_length_ndc, principal_point_ndc.21        img_size: output image resolution (square).22 23    Returns:24        (x_pixels, y_pixels): each (N,) array of pixel coordinates.25    """26    R = np.asarray(camera["R"], dtype=np.float64)27    T = np.asarray(camera["T"], dtype=np.float64)28    focal = camera["focal_length_ndc"]29    pp = camera["principal_point_ndc"]30 31    X_cam = points @ R.T + T32    depth = -X_cam[:, 2]33    Z = np.clip(depth, 1e-4, None)34 35    x_ndc = focal[0] * X_cam[:, 0] / Z + pp[0]36    y_ndc = focal[1] * X_cam[:, 1] / Z + pp[1]37 38    x_px = img_size / 2.0 * (1.0 + x_ndc)39    y_px = img_size / 2.0 * (1.0 - y_ndc)40    return x_px, y_px41 42 43def render_projection(44    image: Image.Image, x_px: np.ndarray, y_px: np.ndarray, img_size: int45) -> np.ndarray:46    """Overlay projected points as a heatmap on an image.47 48    Returns:49        (H, W, 3) float32 blended image in [0, 1].50    """51    valid = (x_px >= 0) & (x_px < img_size) & (y_px >= 0) & (y_px < img_size)52    mask = np.zeros((img_size, img_size), dtype=np.float32)53    mask[y_px[valid].astype(int), x_px[valid].astype(int)] = 1.054    mask = np.clip(gaussian_filter(mask, sigma=1.5) * 5.0, 0, 1)55 56    img_np = np.array(image.convert("RGB")).astype(np.float32) / 255.057    overlay = np.array([1.0, 0.3, 0.1])58    alpha = 0.4 * mask[..., None]59    blended = img_np * (1 - alpha) + overlay * alpha60    return np.clip(blended, 0, 1)61 62 63if __name__ == "__main__":64    import argparse65 66    parser = argparse.ArgumentParser(description="Project point cloud onto an image.")67    parser.add_argument("--image", required=True, help="Path to input image")68    parser.add_argument("--points", required=True, help="Path to surfaces.npy (T,V,6)")69    parser.add_argument(70        "-t",71        "--timestep",72        type=int,73        default=0,74        help="Keyframe index to project (default: 0)",75    )76    parser.add_argument("--camera", required=True, help="Path to camera.json")77    parser.add_argument("--output", required=True, help="Path to output image")78    args = parser.parse_args()79 80    # (T, V, 6) -> take keyframe t, xyz only81    surfaces = np.load(args.points)82    points = surfaces[args.timestep, :, :3].astype(np.float64)83 84    with open(args.camera) as f:85        camera = json.load(f)86    image = Image.open(args.image)87    img_size = image.size[0]88 89    x_px, y_px = project_points(points, camera, img_size)90    result = render_projection(image, x_px, y_px, img_size)91 92    out = Image.fromarray((result * 255).astype(np.uint8))93    out.save(args.output)94    print(f"Saved projection to {args.output}")95