refactor: Enhance play notebook and camera module with undistortion functionality

- Updated the play notebook to reset execution counts and remove unnecessary output displays for clarity.
- Introduced a new `undistort_points` function in the camera module to handle point undistortion using OpenCV.
- Added a `keypoints_undistorted` property to the `Detection` class for lazy evaluation of undistorted keypoints.
- Improved documentation to clarify the usage of keypoints and their undistorted counterparts.

app/camera/__init__.py

@@ -7,11 +7,28 @@ from beartype import beartype
 import jax
 from jax import numpy as jnp
 from jaxtyping import Num, jaxtyped, Array
-from typing_extensions import NotRequired
+from cv2 import undistortPoints
+import numpy as np
 
+NDArray: TypeAlias = np.ndarray
 CameraID: TypeAlias = str  # pylint: disable=invalid-name
 
 
+@jaxtyped(typechecker=beartype)
+def undistort_points(
+    points: Num[NDArray, "M 2"],
+    camera_matrix: Num[NDArray, "3 3"],
+    dist_coeffs: Num[NDArray, "N"],
+) -> Num[NDArray, "M 2"]:
+    """
+    a thin wrapper of cv2.undistortPoints
+    """
+    K = camera_matrix
+    dist = dist_coeffs
+    res = undistortPoints(points, K, dist, P=K)  # type: ignore
+    return res.reshape(-1, 2)
+
+
 @jax.jit
 @jaxtyped(typechecker=beartype)
 def distortion(
@@ -285,7 +302,9 @@ class Detection:
 
     keypoints: Num[Array, "N 2"]
     """
-    Keypoints
+    Keypoints in pixel coordinates. (with camera distortion)
+
+    Use `keypoints_undistorted` to get undistorted keypoints.
     """
     confidences: Num[Array, "N"]
     """
@@ -300,6 +319,24 @@ class Detection:
     Timestamp of the detection
     """
 
+    @property
+    def keypoints_undistorted(self) -> Num[Array, "N 2"]:
+        """
+        Returns undistorted keypoints.
+
+        The result is cached on first access. (lazy evaluation)
+        """
+        kpu = getattr(self, "_kp_undistorted", None)
+        if kpu is None:
+            kpu_np = undistort_points(
+                np.asarray(self.keypoints),
+                np.asarray(self.camera.params.K),
+                np.asarray(self.camera.params.dist_coeffs),
+            )
+            kpu = jnp.asarray(kpu_np)
+            object.__setattr__(self, "_kp_undistorted", kpu)
+        return kpu
+
 
 def classify_by_camera(
     detections: list[Detection],