!pip install visu3d etils[ecolab] jax[cpu] tf-nightly tfds-nightly sunds

from __future__ import annotations
from etils.ecolab.lazy_imports import *

"""
Explanation: Visu3d - Transform (go/v3d-transform)
If you're new to v3d, please look at the intro first.
Installation
We use same installation/imports as in the intro.
End of explanation
"""


tr = v3d.Transform(
    R=[  # Define a rigid rotation
       [-1/3, -(1/3)**.5, (1/3)**.5],
       [1/3, -(1/3)**.5, -(1/3)**.5],
       [-2/3, 0, -(1/3)**.5],
    ],
    t=[2, 2, 2],
)

# Fig display the (x, y, z) basis of the transformation
tr.fig

"""
Explanation: Transformations
v3d makes it easy to project back and forth across coordinates frames.
3d <> 3d
v3d.Transform stores the position, rotation and scale of an object.
It is used to transform objects (e.g. from world to camera 3d coordinates).
v3d.Transform is composed of R (rotation, scale) and t (translation) component:
End of explanation
"""


v3d.make_fig([
    tr,
    tr @ np.array([[0, 0, 0], [1, 1, 1]]),
    tr @ v3d.Point3d(p=[0, 0, 2], rgb=[255, 0, 0]),
    tr @ v3d.Ray(pos=[0, 0, 0], dir=[0, 1, 1]),
    tr @ v3d.Transform(R=np.eye(3), t=[0, 0, 3]),
])

"""
Explanation: v3d.Transform can be composed with all types of objects:

xnp.array
v3d.Ray
v3d.Point3d
v3d.Camera
v3d.Transform
You custom object (see Protocol section below)

Transformation is applied through Python __matmul__ operator: tr @ &lt;obj&gt;
End of explanation
"""


tr.inv

tr.inv @ tr  # `tr.inv @ tr` is identity

"""
Explanation: Inverting a transformation is trivial:
End of explanation
"""


# Camera looking at the center
cam = v3d.Camera.from_look_at(
    spec=v3d.PinholeCamera.from_focal(
        resolution=(128, 170),
        focal_in_px=120,
    ),
    pos=[2, -0.5, 1.7],
    target=[0, 0, 0],  # < TODO(epot): Rename end -> look_at
)

# Point cloud of arbitrary `(..., 3)` shape
rng = np.random.default_rng(0)
point_cloud = v3d.Point3d(
    p=(rng.random((50, 50, 3)) - 0.5) * 3,
    rgb=rng.integers(255, size=(50, 50, 3)),
)

"""
Explanation: See the API for all properties (.matrix4x4, .x_dir, .y_dir, .z_dir,...).
3d <> 2d (Camera pixel projections)
Let's create a camera looking at the center.
End of explanation
"""


# Convert (world 3d) -> (px 2d) coordinates
px_coord = cam.px_from_world @ point_cloud

"""
Explanation: We can project the 3d into 2d pixel coordinates using px_from_world. It supports:

xnp.array: (..., 3) -&gt; (..., 2)
v3d.Point3d -> v3d.Point2d
Your custom objects (see Protocol section below)
End of explanation
"""


# Convert (world 3d) -> (camera 3d) -> (px 2d) coordinates
px_coord = cam.spec.px_from_cam @ cam.cam_from_world @ point_cloud

"""
Explanation: Which is equivalent to:
End of explanation
"""


# Truncate coordinates outside the screen
# Use `(w, h)` as pixels are in `(i, j)` coordinates
px_coord = px_coord.clip(min=0, max=cam.wh)

px_coord.fig

"""
Explanation: v3d.Point2d can be visualized in the pixel space:
End of explanation
"""


px_coord.flatten()[0]

"""
Explanation: v3d.Point3d -> v3d.Point2d will preserve the depth and rgb values, which allows to project back to 3d without any information loss:
End of explanation
"""


print(f'{point_cloud.p.shape} -> {px_coord.p.shape}')

"""
Explanation: The transformation preserves the shape (*shape, 3) -> (*shape, 2).
End of explanation
"""


px_coord = px_coord.replace(depth=None)

# Convert (px 2d) -> (world 3d) coordinates
projected_points = cam.world_from_px @ px_coord

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

"""
Explanation: When the depth is missing, z=1 in camera coordinates:
End of explanation
"""


from etils.array_types import f32

@dataclasses.dataclass(frozen=True)
class MyRay(v3d.DataclassArray):
  pos: f32['*shape 3']
  dir: f32['*shape 3']

  def apply_transform(self, tr: v3d.Transform):
    """Supports `tr @ my_ray`."""
    return self.replace(
        pos=tr @ self.pos,
        # `tr.apply_to_dir` only apply the rotation (tr.R), but NOT the
        # translation (tr.t)
        dir=tr.apply_to_dir(self.dir),
    )


my_ray = MyRay(pos=[0, 0, 0], dir=[0, 0, 1])
cam.world_from_cam @ my_ray

"""
Explanation: Supporting the Transform protocol
To support v3d.Transform, you only need to implement the apply_transform protocol.
End of explanation
"""


