Download lidar frame (#386)

Author: jean-lucas

CHANGELOG.md

@@ -5,6 +5,19 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.15.7](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.7) - 2023-06-09
+
+### Added
+- Allow for downloading pointcloud data for a give task and frame number, example:
+
+```python
+import nucleus
+import numpy as np
+client = nucleus.NucleusClient(API_KEY)
+pts = client.download_pointcloud_task(task_id, frame_num=1)
+np_pts = np.array([pt.to_list() for pt in pts])
+```
+
 ## [0.15.6](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.15.6) - 2023-06-03
 
 ### Changed

nucleus/__init__.py

@@ -17,6 +17,7 @@
     "Keypoint",
     "KeypointsAnnotation",
     "KeypointsPrediction",
+    "LidarPoint",
     "LidarScene",
     "LineAnnotation",
     "LinePrediction",
@@ -57,6 +58,7 @@
     CuboidAnnotation,
     Keypoint,
     KeypointsAnnotation,
+    LidarPoint,
     LineAnnotation,
     MultiCategoryAnnotation,
     Point,
@@ -82,6 +84,7 @@
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
+    I_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
     INDEX_CONTINUOUS_ENABLE_KEY,
@@ -97,6 +100,7 @@
     MODEL_TAGS_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
+    POINTS_KEY,
     PREDICTIONS_IGNORED_KEY,
     PREDICTIONS_PROCESSED_KEY,
     REFERENCE_IDS_KEY,
@@ -978,6 +982,36 @@ def delete_model(self, model_id: str) -> dict:
         )
         return response
 
+    def download_pointcloud_task(
+        self, task_id: str, frame_num: int
+    ) -> List[Union[Point3D, LidarPoint]]:
+        """
+        Download the lidar point cloud data for a give task and frame number.
+
+        Parameters:
+            task_id: download point cloud for this particular task
+            frame_num: download point cloud for this particular frame
+
+        Returns:
+            List of Point3D objects
+
+        """
+
+        response = self.make_request(
+            payload={},
+            route=f"task/{task_id}/frame/{frame_num}",
+            requests_command=requests.get,
+        )
+        points = response.get(POINTS_KEY, None)
+        if points is None or len(points) == 0:
+            raise Exception("Response has invalid payload")
+
+        sample_point = points[0]
+        if I_KEY in sample_point.keys():
+            return [LidarPoint.from_json(pt) for pt in points[:10]]
+
+        return [Point3D.from_json(pt) for pt in points]
+
     @deprecated("Prefer calling Dataset.create_custom_index instead.")
     def create_custom_index(
         self, dataset_id: str, embeddings_urls: list, embedding_dim: int

nucleus/annotation.py

@@ -5,6 +5,8 @@
 from typing import Dict, List, Optional, Sequence, Type, Union
 from urllib.parse import urlparse
 
+import numpy as np
+
 from .constants import (
     ANNOTATION_ID_KEY,
     ANNOTATIONS_KEY,
@@ -15,6 +17,7 @@
     EMBEDDING_VECTOR_KEY,
     GEOMETRY_KEY,
     HEIGHT_KEY,
+    I_KEY,
     INDEX_KEY,
     KEYPOINTS_KEY,
     KEYPOINTS_NAMES_KEY,
@@ -588,6 +591,38 @@ def from_json(cls, payload: Dict[str, float]):
     def to_payload(self) -> dict:
         return {X_KEY: self.x, Y_KEY: self.y, Z_KEY: self.z}
 
+    def to_list(self):
+        return [self.x, self.y, self.z]
+
+
+@dataclass
+class LidarPoint(Point3D):
+    """A Lidar point in 3D space and intensity.
+
+    Parameters:
+        x (float): The x coordinate of the point.
+        y (float): The y coordinate of the point.
+        z (float): The z coordinate of the point.
+        i (float): The intensity value returned by the lidar scan point.
+    """
+
+    i: float
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, float]):
+        return cls(
+            payload[X_KEY], payload[Y_KEY], payload[Z_KEY], payload[I_KEY]
+        )
+
+    def to_payload(self) -> dict:
+        return {X_KEY: self.x, Y_KEY: self.y, Z_KEY: self.z, I_KEY: self.i}
+
+    def to_list(self):
+        return [self.x, self.y, self.z, self.i]
+
+    def to_numpy(self):
+        return np.array(self.to_list())
+
 
 @dataclass  # pylint: disable=R0902
 class CuboidAnnotation(Annotation):  # pylint: disable=R0902

nucleus/constants.py

@@ -108,6 +108,7 @@
 OBJECT_IDS_KEY = "object_ids"
 P1_KEY = "p1"
 P2_KEY = "p2"
+POINTS_KEY = "points"
 POINTCLOUD_KEY = "pointcloud"
 POINTCLOUD_LOCATION_KEY = "pointcloud_location"
 POINTCLOUD_URL_KEY = "pointcloud_url"
@@ -144,6 +145,7 @@
 VISIBLE_KEY = "visible"
 WIDTH_KEY = "width"
 YAW_KEY = "yaw"
+I_KEY = "i"
 W_KEY = "w"
 X_KEY = "x"
 Y_KEY = "y"

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.15.6"
+version = "0.15.7"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <[email protected]>"]