Fit plane tutorial

This tutorial use shows how to generate a plane based on a mesh. Using the Hyperplane and Hyperplane from kornia.gemetry.plane. As data structure we use kornia.geometry.liegroup.So3 e kornia.geometry.vector.Vector3

import plotly.express as px
import plotly.io as pio
import torch
from kornia.core import stack
from kornia.geometry.liegroup import So3
from kornia.geometry.plane import Hyperplane, fit_plane
from kornia.geometry.vector import Vector3
from kornia.utils import create_meshgrid

# define the plane
plane_h = 25
plane_w = 50

# create a base mesh in the ground z == 0
mesh = create_meshgrid(plane_h, plane_w, normalized_coordinates=True)
X, Y = mesh[..., 0], mesh[..., 1]
Z = 0 * X

mesh_pts = Vector3.from_coords(X, Y, Z)

# add noise to the mesh
rand_pts = Vector3.random((plane_h, plane_w))
rand_pts.z.clamp_(min=-0.1, max=0.1)

mesh_view: Vector3 = mesh_pts + rand_pts

x_view = mesh_view.x.ravel().detach().cpu().numpy().tolist()
y_view = mesh_view.y.ravel().detach().cpu().numpy().tolist()
z_view = mesh_view.z.ravel().detach().cpu().numpy().tolist()
fig = px.scatter_3d(dict(x=x_view, y=y_view, z=z_view, category=["view"] * len(x_view)), "x", "y", "z", color="category")
fig.show()

# create rotation
angle_rad = torch.tensor(3.141616 / 4)
rot_x = So3.rot_x(angle_rad)
rot_z = So3.rot_z(angle_rad)
rot = rot_x * rot_z
print(rot)

Parameter containing:
tensor([ 0.8536,  0.3536, -0.1464,  0.3536], requires_grad=True)

# apply the rotation to the mesh points
# TODO: this should work as `rot * mesh_view`
points_rot = stack([rot * x for x in mesh_view.view(-1, 3)]).detach()
points_rot = Vector3(points_rot)

x_rot = points_rot.x.ravel().detach().cpu().numpy().tolist()
y_rot = points_rot.y.ravel().detach().cpu().numpy().tolist()
z_rot = points_rot.z.ravel().detach().cpu().numpy().tolist()

fig = px.scatter_3d(
    dict(
        x=x_view + x_rot,
        y=y_view + y_rot,
        z=z_view + z_rot,
        category=["view"] * len(x_view) + ["rotated"] * len(x_rot),
    ),
    "x",
    "y",
    "z",
    color="category",
)
fig.show()

# estimate the plane from the rotated points
plane_in_ground_fit: Hyperplane = fit_plane(points_rot)
print(plane_in_ground_fit)

Normal: x: -0.0002799616486299783
y: 0.7073221206665039
z: -0.7068911790847778
Offset: 0.094654381275177

# project the original points to the estimated plane
points_proj: Vector3 = plane_in_ground_fit.projection(mesh_view.view(-1, 3))

x_proj = points_proj.x.ravel().detach().cpu().numpy().tolist()
y_proj = points_proj.y.ravel().detach().cpu().numpy().tolist()
z_proj = points_proj.z.ravel().detach().cpu().numpy().tolist()
categories = ["view"] * len(x_view) + ["rotated"] * len(x_rot) + ["projection"] * len(x_proj)
fig = px.scatter_3d(
    dict(
        x=x_view + x_rot + x_proj,
        y=y_view + y_rot + y_proj,
        z=z_view + z_rot + z_proj,
        category=categories,
    ),
    "x",
    "y",
    "z",
    color="category",
)
fig.show()