From 736ba6d1f4ce5b41ef70ff8dc4abd332407ed35f Mon Sep 17 00:00:00 2001
From: John Lambert <johnlambert@gatech.edu>
Date: Wed, 30 Jun 2021 16:25:44 -0400
Subject: [PATCH] use argoverse.utils.transform's yaw_to_quaternion3d instead
 of `quaternion` module

---
 argoverse/evaluation/competition_util.py | 31 +++++++++++++++---------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/argoverse/evaluation/competition_util.py b/argoverse/evaluation/competition_util.py
index 23c64f9d..f7e4fb86 100644
--- a/argoverse/evaluation/competition_util.py
+++ b/argoverse/evaluation/competition_util.py
@@ -10,14 +10,14 @@
 
 import h5py
 import numpy as np
-import quaternion
 from scipy.spatial import ConvexHull
 from shapely.geometry import Polygon
-from sklearn.cluster.dbscan_ import DBSCAN
+from sklearn.cluster import DBSCAN
 
 from argoverse.data_loading.argoverse_tracking_loader import ArgoverseTrackingLoader
 from argoverse.data_loading.object_label_record import ObjectLabelRecord
 from argoverse.utils.se3 import SE3
+from argoverse.utils.transform import yaw_to_quaternion3d
 
 TYPE_LIST = Union[List[np.ndarray], np.ndarray]
 
@@ -169,7 +169,15 @@ def dist(p1: Tuple[float, float], p2: Tuple[float, float]) -> float:
 
 
 def poly_to_label(poly: Polygon, category: str = "VEHICLE", track_id: str = "") -> ObjectLabelRecord:
-    # poly in polygon format
+    """Convert a Shapely Polygon to a 3d cuboid by estimating the minimum-bounding rectangle.
+
+    Args:
+        poly: Shapely polygon object representing a convex hull of an object
+        category: object category to which object belongs, e.g. VEHICLE, PEDESTRIAN, etc
+        track_id: unique identifier
+    Returns:
+        object representing a 3d cuboid
+    """
 
     bbox = poly.minimum_rotated_rectangle
 
@@ -182,6 +190,7 @@ def poly_to_label(poly: Polygon, category: str = "VEHICLE", track_id: str = "")
     d1 = dist((x[0], y[0]), (x[1], y[1]))
     d2 = dist((x[1], y[1]), (x[2], y[2]))
 
+    # assign orientation so that the rectangle's longest side represents the object's length
     width = min(d1, d2)
     length = max(d1, d2)
 
@@ -190,7 +199,8 @@ def poly_to_label(poly: Polygon, category: str = "VEHICLE", track_id: str = "")
     else:
         unit_v = unit_vector((x[0], y[0]), (x[1], y[1]))
 
-    angle = math.atan2(unit_v[1], unit_v[0])
+    angle_rad = np.arctan2(unit_v[1], unit_v[0])
+    q = yaw_to_quaternion3d(angle_rad)
 
     height = max(z) - min(z)
 
@@ -205,10 +215,11 @@ def poly_to_label(poly: Polygon, category: str = "VEHICLE", track_id: str = "")
         ]
     )
 
-    q = quaternion.from_rotation_matrix(R)
+    # location of object in egovehicle coordinates
+    center = np.array([centroid[0], centroid[1], min(z) + height / 2])
 
     return ObjectLabelRecord(
-        quaternion=quaternion.as_float_array(q),
+        quaternion=q,
         translation=center,
         length=length,
         width=width,
@@ -266,18 +277,14 @@ def save_label(argoverse_data: ArgoverseTrackingLoader, labels: List[ObjectLabel
     timestamp = argoverse_data.lidar_timestamp_list[idx]
 
     for label in labels:
+        qw, qx, qy, qz = label.quaternion
         json_data = {
             "center": {
                 "x": label.translation[0],
                 "y": label.translation[1],
                 "z": label.translation[2],
             },
-            "rotation": {
-                "x": label.quaternion[0],
-                "y": label.quaternion[1],
-                "z": label.quaternion[2],
-                "w": label.quaternion[3],
-            },
+            "rotation": {"x": qx, "y": qy, "z": qz, "w": qw},
             "length": label.length,
             "width": label.width,
             "height": label.height,