refactor SceneDatasetItem to DatasetItem

Author: Drew Kaul

nucleus/__init__.py

@@ -109,7 +109,7 @@
     UPDATE_KEY,
 )
 from .dataset import Dataset
-from .dataset_item import DatasetItem
+from .dataset_item import DatasetItem, CameraParams
 from .errors import (
     DatasetItemRetrievalError,
     ModelCreationError,
@@ -135,7 +135,7 @@
 )
 from .slice import Slice
 from .upload_response import UploadResponse
-from .scene import Scene, LidarScene, Frame, SceneDatasetItem, CameraParams
+from .scene import Frame, Scene, LidarScene
 
 # pylint: disable=E1101
 # TODO: refactor to reduce this file to under 1000 lines.

nucleus/annotation.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, List, Optional, Sequence, Union
-from nucleus.dataset_item import is_local_path
+from urllib.parse import urlparse
 
 from .constants import (
     ANNOTATION_ID_KEY,
@@ -310,6 +310,10 @@ def to_payload(self) -> dict:
         }
 
 
+def is_local_path(path: str) -> bool:
+    return urlparse(path).scheme not in {"https", "http", "s3", "gs"}
+
+
 def check_all_mask_paths_remote(
     annotations: Sequence[Union[Annotation]],
 ):

nucleus/constants.py

@@ -64,6 +64,8 @@
 NEW_ITEMS = "new_items"
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+POINTCLOUD_LOCATION_KEY = "pointcloud_location"
+POINTCLOUD_URL_KEY = "pointcloud_url"
 POSITION_KEY = "position"
 PREDICTIONS_IGNORED_KEY = "predictions_ignored"
 PREDICTIONS_PROCESSED_KEY = "predictions_processed"

nucleus/dataset_item.py

@@ -2,36 +2,135 @@
 import json
 import os.path
 from dataclasses import dataclass
-from typing import Optional, Sequence
-from urllib.parse import urlparse
+from typing import Optional, Sequence, Dict, Any
+from enum import Enum
 
+from .annotation import is_local_path, Point3D
 from .constants import (
     DATASET_ITEM_ID_KEY,
     IMAGE_URL_KEY,
     METADATA_KEY,
     ORIGINAL_IMAGE_URL_KEY,
     REFERENCE_ID_KEY,
+    TYPE_KEY,
+    URL_KEY,
+    CAMERA_PARAMS_KEY,
+    POINTCLOUD_URL_KEY,
+    X_KEY,
+    Y_KEY,
+    Z_KEY,
+    W_KEY,
+    POSITION_KEY,
+    HEADING_KEY,
+    FX_KEY,
+    FY_KEY,
+    CX_KEY,
+    CY_KEY,
 )
 
 
 @dataclass
-class DatasetItem:
+class Quaternion:
+    x: float
+    y: float
+    z: float
+    w: float
 
-    image_location: str
+    @classmethod
+    def from_json(cls, payload: Dict[str, float]):
+        return cls(
+            payload[X_KEY], payload[Y_KEY], payload[Z_KEY], payload[W_KEY]
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            X_KEY: self.x,
+            Y_KEY: self.y,
+            Z_KEY: self.z,
+            W_KEY: self.w,
+        }
+
+
+@dataclass
+class CameraParams:
+    position: Point3D
+    heading: Quaternion
+    fx: float
+    fy: float
+    cx: float
+    cy: float
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, Any]):
+        return cls(
+            Point3D.from_json(payload[POSITION_KEY]),
+            Quaternion.from_json(payload[HEADING_KEY]),
+            payload[FX_KEY],
+            payload[FY_KEY],
+            payload[CX_KEY],
+            payload[CY_KEY],
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            POSITION_KEY: self.position.to_payload(),
+            HEADING_KEY: self.heading.to_payload(),
+            FX_KEY: self.fx,
+            FY_KEY: self.fy,
+            CX_KEY: self.cx,
+            CY_KEY: self.cy,
+        }
+
+
+class DatasetItemType(Enum):
+    IMAGE = "image"
+    POINTCLOUD = "pointcloud"
+
+
+@dataclass  # pylint: disable=R0902
+class DatasetItem:  # pylint: disable=R0902
+    image_location: Optional[str] = None
     reference_id: Optional[str] = None
     item_id: Optional[str] = None
     metadata: Optional[dict] = None
+    pointcloud_location: Optional[str] = None
 
     def __post_init__(self):
         self.local = is_local_path(self.image_location)
+        assert bool(self.image_location) != bool(
+            self.pointcloud_location
+        ), "Must specify exactly one of the image_location, pointcloud_location parameters"
+        self.type = (
+            DatasetItemType.IMAGE
+            if self.image_location
+            else DatasetItemType.POINTCLOUD
+        )
+        camera_params = (
+            self.metadata.get(CAMERA_PARAMS_KEY, None)
+            if self.metadata
+            else None
+        )
+        self.camera_params = (
+            CameraParams.from_json(camera_params) if camera_params else None
+        )
 
     @classmethod
-    def from_json(cls, payload: dict):
-        url = payload.get(IMAGE_URL_KEY, "") or payload.get(
+    def from_json(cls, payload: dict, is_scene=False):
+        image_url = payload.get(IMAGE_URL_KEY, "") or payload.get(
             ORIGINAL_IMAGE_URL_KEY, ""
         )
+
+        if is_scene:
+            return cls(
+                image_location=image_url,
+                pointcloud_location=payload.get(POINTCLOUD_URL_KEY, ""),
+                reference_id=payload.get(REFERENCE_ID_KEY, None),
+                item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+                metadata=payload.get(METADATA_KEY, {}),
+            )
+
         return cls(
-            image_location=url,
+            image_location=image_url,
             reference_id=payload.get(REFERENCE_ID_KEY, None),
             item_id=payload.get(DATASET_ITEM_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
@@ -40,28 +139,39 @@ def from_json(cls, payload: dict):
     def local_file_exists(self):
         return os.path.isfile(self.image_location)
 
-    def to_payload(self) -> dict:
-        payload = {
-            IMAGE_URL_KEY: self.image_location,
+    def to_payload(self, is_scene=False) -> dict:
+        payload: Dict[str, Any] = {
             METADATA_KEY: self.metadata or {},
         }
+
+        if is_scene:
+            if self.image_location:
+                payload[URL_KEY] = self.image_location
+            elif self.pointcloud_location:
+                payload[URL_KEY] = self.pointcloud_location
+            payload[TYPE_KEY] = self.type.value
+        else:
+            assert (
+                self.image_location
+            ), "Must specify image_location for DatasetItems not in a Scene"
+            payload[IMAGE_URL_KEY] = self.image_location
+
         if self.reference_id:
             payload[REFERENCE_ID_KEY] = self.reference_id
         if self.item_id:
             payload[DATASET_ITEM_ID_KEY] = self.item_id
+        if self.camera_params:
+            payload[CAMERA_PARAMS_KEY] = self.camera_params
+
         return payload
 
     def to_json(self) -> str:
         return json.dumps(self.to_payload(), allow_nan=False)
 
 
-def is_local_path(path: str) -> bool:
-    return urlparse(path).scheme not in {"https", "http", "s3", "gs"}
-
-
 def check_all_paths_remote(dataset_items: Sequence[DatasetItem]):
     for item in dataset_items:
-        if is_local_path(item.image_location):
+        if item.image_location and is_local_path(item.image_location):
             raise ValueError(
                 f"All paths must be remote, but {item.image_location} is either "
                 "local, or a remote URL type that is not supported."