fix to_payload fn's

Author: Drew Kaul

nucleus/constants.py

@@ -15,6 +15,8 @@
 CAMERA_PARAMS_KEY = "camera_params"
 CLASS_PDF_KEY = "class_pdf"
 CONFIDENCE_KEY = "confidence"
+CX_KEY = "cx"
+CY_KEY = "cy"
 DATASET_ID_KEY = "dataset_id"
 DATASET_ITEM_IDS_KEY = "dataset_item_ids"
 DATASET_ITEM_ID_KEY = "dataset_item_id"
@@ -32,7 +34,10 @@
 ERROR_ITEMS = "upload_errors"
 ERROR_PAYLOAD = "error_payload"
 FRAMES_KEY = "frames"
+FX_KEY = "fx"
+FY_KEY = "fy"
 GEOMETRY_KEY = "geometry"
+HEADING_KEY = "heading"
 HEIGHT_KEY = "height"
 IGNORED_ITEMS = "ignored_items"
 IMAGE_KEY = "image"
@@ -77,6 +82,7 @@
 VERTICES_KEY = "vertices"
 WIDTH_KEY = "width"
 YAW_KEY = "yaw"
+W_KEY = "w"
 X_KEY = "x"
 Y_KEY = "y"
 Z_KEY = "z"

nucleus/scene.py

@@ -1,17 +1,27 @@
 import json
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import Optional, Union, Dict, List, Set
+from typing import Optional, Union, Any, Dict, List, Set
 from enum import Enum
 from nucleus.constants import (
     CAMERA_PARAMS_KEY,
+    CX_KEY,
+    CY_KEY,
     FRAMES_KEY,
+    FX_KEY,
+    FY_KEY,
+    HEADING_KEY,
     INDEX_KEY,
     ITEMS_KEY,
     METADATA_KEY,
+    POSITION_KEY,
     REFERENCE_ID_KEY,
     TYPE_KEY,
     URL_KEY,
+    W_KEY,
+    X_KEY,
+    Y_KEY,
+    Z_KEY,
 )
 from .annotation import Point3D
 from .utils import flatten
@@ -29,6 +39,20 @@ class Quaternion:
     z: float
     w: float
 
+    @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:
@@ -39,6 +63,27 @@ class CameraParams:
     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,
+        }
+
 
 @dataclass
 class SceneDatasetItem:
@@ -60,17 +105,21 @@ def from_json(cls, payload: dict):
             type=payload.get(TYPE_KEY, ""),
             reference_id=payload.get(REFERENCE_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, None),
-            camera_params=payload.get(CAMERA_PARAMS_KEY, None),
+            camera_params=CameraParams.from_json(
+                payload.get(CAMERA_PARAMS_KEY, {})
+            ),
         )
 
     def to_payload(self) -> dict:
-        return {
+        payload = {
             URL_KEY: self.url,
             TYPE_KEY: self.type,
             REFERENCE_ID_KEY: self.reference_id,
             METADATA_KEY: self.metadata,
-            CAMERA_PARAMS_KEY: self.camera_params,
         }
+        if self.camera_params:
+            payload[CAMERA_PARAMS_KEY] = self.camera_params.to_payload()
+        return payload
 
     def to_json(self) -> str:
         return json.dumps(self.to_payload(), allow_nan=False)