10000 GitHub - google-research/visu3d: 3d without friction (Torch, TF, Jax, Numpy)
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

google-research/visu3d

Repository files navigation

visu3d - 3D geometry made easy

Unittests PyPI version Documentation Status

visu3d is an abstraction layer between Torch/TF/Jax/Numpy and your program. It provides:

  • Standard primitives for 3d geometry (Ray, Camera, Transform,...). You can combine those standard primitives with your custom ones.
  • Everything is trivially visualizable with zero boilerplate. Inspect & debug camera poses, trajectories,...
  • All primitives are dataclass_array, dataclasses which can be reshaped, sliced,... as if they were numpy arrays.
  • Everything is extensible, you can gradually opt-in only for the features you need, and replace any standard primitive by your custom ones.

Core features

Everything is a v3d.DataclassArray: dataclass behave like np.array (with indexing, slicing, shape manipulation, vectorization,...).

# Single ray
ray = v3d.Ray(pos=[0, 0, 0], dir=[1, 1, 1])
assert rays.shape == ()

# Multiple rays batched together
rays = v3d.Ray(pos=np.zeros((B, H, W, 3)), dir=np.ones((B, H, W, 3)))
assert rays.shape == (B, H, W)

rays = rays.reshape('b h w -> b (h w)')  #  Native `einops` support

top_left_ray = rays[..., 0, 0]  #  (b, h, w) -> (b,)

rays = rays.flatten()
rays = rays[rays.norm() > 0]  # Filter invalid rays

Everything is visualizable interactively

Every object has a .fig property for interactive visualization:

rays.fig  # Plot the rays

Display multiple objects together:

v3d.make_fig([cam, rays, point_cloud])

Auto-plot figures with Colab magic:

v3d.auto_plot_figs()  # Once at the start of the Colab

cam, rays, point_cloud  # Tuple auto-displayed without `v3d.make_fig` call

Same code seamlessly works across Torch, Jax, TensorFlow, Numpy.

rays = rays.as_jax()  # .as_tf(), as_np(), .as_jax()
assert isinstance(rays.pos, jnp.ndarray)
assert rays.xnp is jnp

rays = rays.as_torch()
assert isinstance(rays.pos, torch.Tensor)
assert rays.xnp is torch

With native support for auto-diff, jax.vmap, jax.tree_utils,...

Privitives

Common primitives (Ray, Camera, Transform,...), so user can express intent, instead of math.

H, W = (256, 1024)
cam_spec = v3d.PinholeCamera.from_focal(
    resolution=(H, W),
    focal_in_px=35.,
)
cam = v3d.Camera.from_look_at(
    spec=cam_spec,
    pos=[5, 5, 5],
    target=[0, 0, 0],  # Camera looks at the scene center
)

rays = cam.rays()  # Rays in world coordinates

# Project `(*shape, 3)` world coordinates into `(*shape, 2)` pixel coordinates.
px_coords = cam.px_from_world @ point_cloud

See the API for a full list of primitive.

Creating your own primitives is trivial.

Converting any dataclass to dataclass array is trivial:

  • Inherit from v3d.DataclassArray
  • Use etils.array_types to annotate the array fields (or exlicitly use my_field: Any = dca.field(shape=, dtype=) instead of dataclasses.field)
from etils.array_types import FloatArray


class MyRay(v3d.DataclassArray):
  pos: FloatArray[..., 3]
  dir: FloatArray[..., 3]


rays = MyRay(pos=jnp.zeros((H, W, 3)), dir=jnp.ones((H, W, 3)))
assert rays.shape == (H, W)

v3d makes it easy to opt-in to the feature you need by implementing the corresponding protocol.

Documentation

The best way to get started is to try the Colab tutorials (in the documentation):

Installation:

pip install visu3d

Usage:

import visu3d as v3d

Installation

pip install visu3d

This is not an official Google product.

About

3d without friction (Torch, TF, Jax, Numpy)

Topics

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Contributors 8

Languages

0