From 4ad0afb67a3977fd1690fb1347038dfc93aace62 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Mon, 19 Jun 2023 13:27:42 +0200
Subject: [PATCH 1/9] refactor random_flip implementation

---
 pcdet/datasets/augmentor/augmentor_utils.py | 33 +++++++--------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index 3c088e33c..e28100c7c 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -6,42 +6,31 @@
 
 
 def random_flip_along_x(gt_boxes, points, return_flip=False, enable=None):
-    """
-    Args:
-        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
-        points: (M, 3 + C)
-    Returns:
-    """
-    if enable is None:
-        enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
-    if enable:
-        gt_boxes[:, 1] = -gt_boxes[:, 1]
-        gt_boxes[:, 6] = -gt_boxes[:, 6]
-        points[:, 1] = -points[:, 1]
-        
-        if gt_boxes.shape[1] > 7:
-            gt_boxes[:, 8] = -gt_boxes[:, 8]
-    if return_flip:
-        return gt_boxes, points, enable
-    return gt_boxes, points
+    return random_flip_along(0, gt_boxes, points, return_flip=return_flip, enable=enable)
 
 
 def random_flip_along_y(gt_boxes, points, return_flip=False, enable=None):
+    return random_flip_along(1, gt_boxes, points, return_flip=return_flip, enable=enable)
+
+
+def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
     """
     Args:
         gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C)
     Returns:
     """
+    assert dim in [0, 1]
+    other_dim = 1 - dim
     if enable is None:
         enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
     if enable:
-        gt_boxes[:, 0] = -gt_boxes[:, 0]
-        gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)
-        points[:, 0] = -points[:, 0]
+        gt_boxes[:, other_dim] = -gt_boxes[:, other_dim]
+        gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi * dim)
+        points[:, other_dim] = -points[:, other_dim]
 
         if gt_boxes.shape[1] > 7:
-            gt_boxes[:, 7] = -gt_boxes[:, 7]
+            gt_boxes[:, 7 + other_dim] = -gt_boxes[:, 7 + other_dim]
     if return_flip:
         return gt_boxes, points, enable
     return gt_boxes, points

From 54d755b5a3c3c6f0fcef21fd6a2168b7bd0b34e3 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Mon, 19 Jun 2023 23:12:08 +0200
Subject: [PATCH 2/9] refactor random_flip usage

---
 pcdet/datasets/augmentor/augmentor_utils.py | 10 +---------
 pcdet/datasets/augmentor/data_augmentor.py  | 12 ++++++++----
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index e28100c7c..6d019d978 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -5,14 +5,6 @@
 from ...utils import box_utils
 
 
-def random_flip_along_x(gt_boxes, points, return_flip=False, enable=None):
-    return random_flip_along(0, gt_boxes, points, return_flip=return_flip, enable=enable)
-
-
-def random_flip_along_y(gt_boxes, points, return_flip=False, enable=None):
-    return random_flip_along(1, gt_boxes, points, return_flip=return_flip, enable=enable)
-
-
 def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
     """
     Args:
@@ -20,7 +12,7 @@ def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
         points: (M, 3 + C)
     Returns:
     """
-    assert dim in [0, 1]
+    assert dim in [0, 1]  # corresponds to x-, y-axis respectively
     other_dim = 1 - dim
     if enable is None:
         enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
diff --git a/pcdet/datasets/augmentor/data_augmentor.py b/pcdet/datasets/augmentor/data_augmentor.py
index 56acebc81..03301ed46 100644
--- a/pcdet/datasets/augmentor/data_augmentor.py
+++ b/pcdet/datasets/augmentor/data_augmentor.py
@@ -59,14 +59,18 @@ def random_world_flip(self, data_dict=None, config=None):
         gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
         for cur_axis in config['ALONG_AXIS_LIST']:
             assert cur_axis in ['x', 'y']
-            gt_boxes, points, enable = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)(
-                gt_boxes, points, return_flip=True
+            cur_dim = ['x', 'y'].index(cur_axis)
+            gt_boxes, points, enable = augmentor_utils.random_flip_along(
+                cur_dim, gt_boxes, points, return_flip=True
             )
             data_dict['flip_%s'%cur_axis] = enable
             if 'roi_boxes' in data_dict.keys():
                 num_frame, num_rois,dim = data_dict['roi_boxes'].shape
-                roi_boxes, _, _ = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)(
-                data_dict['roi_boxes'].reshape(-1,dim), np.zeros([1,3]), return_flip=True, enable=enable
+                roi_boxes, _ = augmentor_utils.random_flip_along(
+                    cur_dim,
+                    data_dict['roi_boxes'].reshape(-1,dim),
+                    np.zeros([0,3]),
+                    enable=enable,
                 )
                 data_dict['roi_boxes'] = roi_boxes.reshape(num_frame, num_rois,dim)
 

From ddfd89ea611d0cf987f55748e47bbe04f64415d5 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Mon, 19 Jun 2023 23:15:05 +0200
Subject: [PATCH 3/9] refactor random flip to be batch dim independent

---
 pcdet/datasets/augmentor/augmentor_utils.py | 14 +++++++-------
 pcdet/datasets/augmentor/data_augmentor.py  |  6 ++----
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index 6d019d978..1984693ce 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -8,8 +8,8 @@
 def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
     """
     Args:
-        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
-        points: (M, 3 + C)
+        gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        points: (*, 3 + C)
     Returns:
     """
     assert dim in [0, 1]  # corresponds to x-, y-axis respectively
@@ -17,12 +17,12 @@ def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
     if enable is None:
         enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
     if enable:
-        gt_boxes[:, other_dim] = -gt_boxes[:, other_dim]
-        gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi * dim)
-        points[:, other_dim] = -points[:, other_dim]
+        gt_boxes[..., other_dim] = -gt_boxes[..., other_dim]
+        gt_boxes[..., 6] = -(gt_boxes[..., 6] + np.pi * dim)
+        points[..., other_dim] = -points[..., other_dim]
 
-        if gt_boxes.shape[1] > 7:
-            gt_boxes[:, 7 + other_dim] = -gt_boxes[:, 7 + other_dim]
+        if gt_boxes.shape[-1] > 7:
+            gt_boxes[..., 7 + other_dim] = -gt_boxes[..., 7 + other_dim]
     if return_flip:
         return gt_boxes, points, enable
     return gt_boxes, points
diff --git a/pcdet/datasets/augmentor/data_augmentor.py b/pcdet/datasets/augmentor/data_augmentor.py
index 03301ed46..7681f402c 100644
--- a/pcdet/datasets/augmentor/data_augmentor.py
+++ b/pcdet/datasets/augmentor/data_augmentor.py
@@ -65,14 +65,12 @@ def random_world_flip(self, data_dict=None, config=None):
             )
             data_dict['flip_%s'%cur_axis] = enable
             if 'roi_boxes' in data_dict.keys():
-                num_frame, num_rois,dim = data_dict['roi_boxes'].shape
-                roi_boxes, _ = augmentor_utils.random_flip_along(
+                data_dict['roi_boxes'], _ = augmentor_utils.random_flip_along(
                     cur_dim,
-                    data_dict['roi_boxes'].reshape(-1,dim),
+                    data_dict['roi_boxes'],
                     np.zeros([0,3]),
                     enable=enable,
                 )
-                data_dict['roi_boxes'] = roi_boxes.reshape(num_frame, num_rois,dim)
 
         data_dict['gt_boxes'] = gt_boxes
         data_dict['points'] = points

From 5c226dc832a22cd896b572b518e1393f526ccc4e Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Mon, 19 Jun 2023 22:57:09 +0200
Subject: [PATCH 4/9] refactor random_world_rotation

---
 pcdet/datasets/augmentor/augmentor_utils.py | 14 +++++++-------
 pcdet/datasets/augmentor/data_augmentor.py  | 10 ++++++----
 pcdet/utils/common_utils.py                 |  7 ++++++-
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index 1984693ce..2b34f2af7 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -31,7 +31,7 @@ def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
 def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotation=None):
     """
     Args:
-        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C),
         rot_range: [min, max]
     Returns:
@@ -39,13 +39,13 @@ def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotatio
     if noise_rotation is None: 
         noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
     points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
-    gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0]
-    gt_boxes[:, 6] += noise_rotation
-    if gt_boxes.shape[1] > 7:
-        gt_boxes[:, 7:9] = common_utils.rotate_points_along_z(
-            np.hstack((gt_boxes[:, 7:9], np.zeros((gt_boxes.shape[0], 1))))[np.newaxis, :, :],
+    gt_boxes[..., 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, ..., 0:3], np.array([noise_rotation]))[0]
+    gt_boxes[..., 6] += noise_rotation
+    if gt_boxes.shape[-1] > 7:
+        gt_boxes[..., 7:9] = common_utils.rotate_points_along_z(
+            np.concatenate((gt_boxes[..., 7:9], np.zeros((*gt_boxes.shape[:-1], 1))), axis=-1)[np.newaxis, ...],
             np.array([noise_rotation])
-        )[0][:, 0:2]
+        )[0, ..., 0:2]
 
     if return_rot:
         return gt_boxes, points, noise_rotation
diff --git a/pcdet/datasets/augmentor/data_augmentor.py b/pcdet/datasets/augmentor/data_augmentor.py
index 7681f402c..0945f1de6 100644
--- a/pcdet/datasets/augmentor/data_augmentor.py
+++ b/pcdet/datasets/augmentor/data_augmentor.py
@@ -86,10 +86,12 @@ def random_world_rotation(self, data_dict=None, config=None):
             data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range, return_rot=True
         )
         if 'roi_boxes' in data_dict.keys():
-            num_frame, num_rois,dim = data_dict['roi_boxes'].shape
-            roi_boxes, _, _ = augmentor_utils.global_rotation(
-            data_dict['roi_boxes'].reshape(-1, dim), np.zeros([1, 3]), rot_range=rot_range, return_rot=True, noise_rotation=noise_rot)
-            data_dict['roi_boxes'] = roi_boxes.reshape(num_frame, num_rois,dim)
+            data_dict['roi_boxes'], _ = augmentor_utils.global_rotation(
+                data_dict['roi_boxes'],
+                np.zeros([0, 3]),
+                rot_range=rot_range,
+                noise_rotation=noise_rot,
+            )
 
         data_dict['gt_boxes'] = gt_boxes
         data_dict['points'] = points
diff --git a/pcdet/utils/common_utils.py b/pcdet/utils/common_utils.py
index af70728db..8864ceb46 100644
--- a/pcdet/utils/common_utils.py
+++ b/pcdet/utils/common_utils.py
@@ -35,7 +35,7 @@ def drop_info_with_name(info, name):
 def rotate_points_along_z(points, angle):
     """
     Args:
-        points: (B, N, 3 + C)
+        points: (B, *, 3 + C)
         angle: (B), angle along z-axis, angle increases x ==> y
     Returns:
 
@@ -43,6 +43,10 @@ def rotate_points_along_z(points, angle):
     points, is_numpy = check_numpy_to_torch(points)
     angle, _ = check_numpy_to_torch(angle)
 
+    orig_shape = points.shape
+    if len(orig_shape) > 3:
+        points = points.view(orig_shape[0], -1, orig_shape[-1])
+
     cosa = torch.cos(angle)
     sina = torch.sin(angle)
     zeros = angle.new_zeros(points.shape[0])
@@ -54,6 +58,7 @@ def rotate_points_along_z(points, angle):
     ), dim=1).view(-1, 3, 3).float()
     points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
     points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
+    points_rot = points_rot.view(orig_shape)
     return points_rot.numpy() if is_numpy else points_rot
 
 

From 8a76f5969443b064d3b39b7b4ea2ae86386928b3 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Fri, 23 Jun 2023 17:36:05 +0200
Subject: [PATCH 5/9] fix scaling with roi boxes

---
 pcdet/datasets/augmentor/augmentor_utils.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index 2b34f2af7..7820905b9 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -55,7 +55,7 @@ def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotatio
 def global_scaling(gt_boxes, points, scale_range, return_scale=False):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C),
         scale_range: [min, max]
     Returns:
@@ -66,7 +66,7 @@ def global_scaling(gt_boxes, points, scale_range, return_scale=False):
     points[:, :3] *= noise_scale
     gt_boxes[:, :6] *= noise_scale
     if gt_boxes.shape[1] > 7:
-        gt_boxes[:, 7:] *= noise_scale
+        gt_boxes[:, 7:9] *= noise_scale
         
     if return_scale:
         return gt_boxes, points, noise_scale
@@ -75,7 +75,7 @@ def global_scaling(gt_boxes, points, scale_range, return_scale=False):
 def global_scaling_with_roi_boxes(gt_boxes, roi_boxes, points, scale_range, return_scale=False):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
+        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C),
         scale_range: [min, max]
     Returns:
@@ -85,7 +85,11 @@ def global_scaling_with_roi_boxes(gt_boxes, roi_boxes, points, scale_range, retu
     noise_scale = np.random.uniform(scale_range[0], scale_range[1])
     points[:, :3] *= noise_scale
     gt_boxes[:, :6] *= noise_scale
+    if gt_boxes.shape[1] > 7:
+        gt_boxes[:, 7:9] *= noise_scale
+
     roi_boxes[:,:, [0,1,2,3,4,5,7,8]] *= noise_scale
+
     if return_scale:
         return gt_boxes,roi_boxes, points, noise_scale
     return gt_boxes, roi_boxes, points

From b4dd91578b3645ae503d55acda1ff691f913134b Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Fri, 23 Jun 2023 17:56:13 +0200
Subject: [PATCH 6/9] fix roi_boxes noise translation

---
 pcdet/datasets/augmentor/data_augmentor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pcdet/datasets/augmentor/data_augmentor.py b/pcdet/datasets/augmentor/data_augmentor.py
index 0945f1de6..fd0fef648 100644
--- a/pcdet/datasets/augmentor/data_augmentor.py
+++ b/pcdet/datasets/augmentor/data_augmentor.py
@@ -152,7 +152,7 @@ def random_world_translation(self, data_dict=None, config=None):
         gt_boxes[:, :3] += noise_translate
                 
         if 'roi_boxes' in data_dict.keys():
-            data_dict['roi_boxes'][:, :3] += noise_translate
+            data_dict['roi_boxes'][:, :, :3] += noise_translate
         
         data_dict['gt_boxes'] = gt_boxes
         data_dict['points'] = points

From 4436d06622a814c0fe4160a07d401f64fe7ed5e1 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Fri, 23 Jun 2023 23:32:32 +0200
Subject: [PATCH 7/9] refactor application of data_dict transforms

---
 pcdet/datasets/augmentor/augmentor_utils.py | 106 +++++++++++---------
 pcdet/datasets/augmentor/data_augmentor.py  |  58 +++--------
 pcdet/datasets/processor/data_processor.py  |  13 +--
 pcdet/utils/common_utils.py                 |  12 +++
 4 files changed, 90 insertions(+), 99 deletions(-)

diff --git a/pcdet/datasets/augmentor/augmentor_utils.py b/pcdet/datasets/augmentor/augmentor_utils.py
index 7820905b9..075b44ba4 100644
--- a/pcdet/datasets/augmentor/augmentor_utils.py
+++ b/pcdet/datasets/augmentor/augmentor_utils.py
@@ -5,7 +5,7 @@
 from ...utils import box_utils
 
 
-def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
+def random_flip_along(dim, return_flip=False, enable=None):
     """
     Args:
         gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
@@ -16,19 +16,31 @@ def random_flip_along(dim, gt_boxes, points, return_flip=False, enable=None):
     other_dim = 1 - dim
     if enable is None:
         enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
+
     if enable:
-        gt_boxes[..., other_dim] = -gt_boxes[..., other_dim]
-        gt_boxes[..., 6] = -(gt_boxes[..., 6] + np.pi * dim)
-        points[..., other_dim] = -points[..., other_dim]
+        def flip_pointlike(points):
+            points[..., other_dim] = -points[..., other_dim]
+            return points
+
+        def flip_boxlike(boxes):
+            boxes[..., other_dim] = -boxes[..., other_dim]
+            boxes[..., 6] = -(boxes[..., 6] + np.pi * dim)
+
+            if boxes.shape[-1] > 7:
+                boxes[..., 7 + other_dim] = -boxes[..., 7 + other_dim]
+
+            return boxes
+
+        tfs = dict(point=flip_pointlike, box=flip_boxlike)
+    else:
+        tfs = dict()
 
-        if gt_boxes.shape[-1] > 7:
-            gt_boxes[..., 7 + other_dim] = -gt_boxes[..., 7 + other_dim]
     if return_flip:
-        return gt_boxes, points, enable
-    return gt_boxes, points
+        return tfs, enable
+    return tfs
 
 
-def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotation=None):
+def global_rotation(rot_range, return_rot=False, noise_rotation=None):
     """
     Args:
         gt_boxes: (*, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
@@ -38,61 +50,59 @@ def global_rotation(gt_boxes, points, rot_range, return_rot=False, noise_rotatio
     """
     if noise_rotation is None: 
         noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
-    points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
-    gt_boxes[..., 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, ..., 0:3], np.array([noise_rotation]))[0]
-    gt_boxes[..., 6] += noise_rotation
-    if gt_boxes.shape[-1] > 7:
-        gt_boxes[..., 7:9] = common_utils.rotate_points_along_z(
-            np.concatenate((gt_boxes[..., 7:9], np.zeros((*gt_boxes.shape[:-1], 1))), axis=-1)[np.newaxis, ...],
-            np.array([noise_rotation])
-        )[0, ..., 0:2]
+
+    def rotate_pointlike(points):
+        points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
+        return points
+
+    def rotate_boxlike(boxes):
+        boxes[..., 0:3] = common_utils.rotate_points_along_z(boxes[np.newaxis, ..., 0:3], np.array([noise_rotation]))[0]
+        boxes[..., 6] += noise_rotation
+        if boxes.shape[-1] > 7:
+            boxes[..., 7:9] = common_utils.rotate_points_along_z(
+                np.concatenate((boxes[..., 7:9], np.zeros((*boxes.shape[:-1], 1))), axis=-1)[np.newaxis, ...],
+                np.array([noise_rotation])
+            )[0, ..., 0:2]
+        return boxes
+
+    tfs = dict(point=rotate_pointlike, box=rotate_boxlike)
 
     if return_rot:
-        return gt_boxes, points, noise_rotation
-    return gt_boxes, points
+        return tfs, noise_rotation
+    return tfs
 
 
-def global_scaling(gt_boxes, points, scale_range, return_scale=False):
+def global_scaling(scale_range, return_scale=False):
     """
     Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
+        gt_boxes: (*, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
         points: (M, 3 + C),
         scale_range: [min, max]
     Returns:
     """
     if scale_range[1] - scale_range[0] < 1e-3:
-        return gt_boxes, points
+        noise_scale = sum(scale_range) / len(scale_range)
+        assert noise_scale == 1.0, (noise_scale, scale_range)
     noise_scale = np.random.uniform(scale_range[0], scale_range[1])
-    points[:, :3] *= noise_scale
-    gt_boxes[:, :6] *= noise_scale
-    if gt_boxes.shape[1] > 7:
-        gt_boxes[:, 7:9] *= noise_scale
-        
-    if return_scale:
-        return gt_boxes, points, noise_scale
-    return gt_boxes, points
 
-def global_scaling_with_roi_boxes(gt_boxes, roi_boxes, points, scale_range, return_scale=False):
-    """
-    Args:
-        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
-        points: (M, 3 + C),
-        scale_range: [min, max]
-    Returns:
-    """
-    if scale_range[1] - scale_range[0] < 1e-3:
-        return gt_boxes, points
-    noise_scale = np.random.uniform(scale_range[0], scale_range[1])
-    points[:, :3] *= noise_scale
-    gt_boxes[:, :6] *= noise_scale
-    if gt_boxes.shape[1] > 7:
-        gt_boxes[:, 7:9] *= noise_scale
+    def scale_pointlike(points):
+        points[:, :3] *= noise_scale
+        return points
+
+    def scale_boxlike(boxes):
+        boxes[..., :6] *= noise_scale
+        if boxes.shape[-1] > 7:
+            boxes[..., 7:9] *= noise_scale
+        return boxes
 
-    roi_boxes[:,:, [0,1,2,3,4,5,7,8]] *= noise_scale
+    if noise_scale != 1.0:
+        tfs = dict(point=scale_pointlike, box=scale_boxlike)
+    else:
+        tfs = {}
 
     if return_scale:
-        return gt_boxes,roi_boxes, points, noise_scale
-    return gt_boxes, roi_boxes, points
+        return tfs, noise_scale
+    return tfs
 
 
 def random_image_flip_horizontal(image, depth_map, gt_boxes, calib):
diff --git a/pcdet/datasets/augmentor/data_augmentor.py b/pcdet/datasets/augmentor/data_augmentor.py
index fd0fef648..2d5f64f18 100644
--- a/pcdet/datasets/augmentor/data_augmentor.py
+++ b/pcdet/datasets/augmentor/data_augmentor.py
@@ -56,24 +56,15 @@ def __setstate__(self, d):
     def random_world_flip(self, data_dict=None, config=None):
         if data_dict is None:
             return partial(self.random_world_flip, config=config)
-        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
         for cur_axis in config['ALONG_AXIS_LIST']:
             assert cur_axis in ['x', 'y']
             cur_dim = ['x', 'y'].index(cur_axis)
-            gt_boxes, points, enable = augmentor_utils.random_flip_along(
-                cur_dim, gt_boxes, points, return_flip=True
+            tfs, enable = augmentor_utils.random_flip_along(
+                cur_dim, return_flip=True
             )
+            common_utils.apply_data_transform(data_dict, tfs)
             data_dict['flip_%s'%cur_axis] = enable
-            if 'roi_boxes' in data_dict.keys():
-                data_dict['roi_boxes'], _ = augmentor_utils.random_flip_along(
-                    cur_dim,
-                    data_dict['roi_boxes'],
-                    np.zeros([0,3]),
-                    enable=enable,
-                )
 
-        data_dict['gt_boxes'] = gt_boxes
-        data_dict['points'] = points
         return data_dict
 
     def random_world_rotation(self, data_dict=None, config=None):
@@ -82,38 +73,18 @@ def random_world_rotation(self, data_dict=None, config=None):
         rot_range = config['WORLD_ROT_ANGLE']
         if not isinstance(rot_range, list):
             rot_range = [-rot_range, rot_range]
-        gt_boxes, points, noise_rot = augmentor_utils.global_rotation(
-            data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range, return_rot=True
+        tfs, noise_rot = augmentor_utils.global_rotation(
+            rot_range=rot_range, return_rot=True
         )
-        if 'roi_boxes' in data_dict.keys():
-            data_dict['roi_boxes'], _ = augmentor_utils.global_rotation(
-                data_dict['roi_boxes'],
-                np.zeros([0, 3]),
-                rot_range=rot_range,
-                noise_rotation=noise_rot,
-            )
-
-        data_dict['gt_boxes'] = gt_boxes
-        data_dict['points'] = points
+        common_utils.apply_data_transform(data_dict, tfs)
         data_dict['noise_rot'] = noise_rot
         return data_dict
 
     def random_world_scaling(self, data_dict=None, config=None):
         if data_dict is None:
             return partial(self.random_world_scaling, config=config)
-        
-        if 'roi_boxes' in data_dict.keys():
-            gt_boxes, roi_boxes, points, noise_scale = augmentor_utils.global_scaling_with_roi_boxes(
-                data_dict['gt_boxes'], data_dict['roi_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE'], return_scale=True
-            )
-            data_dict['roi_boxes'] = roi_boxes
-        else:
-            gt_boxes, points, noise_scale = augmentor_utils.global_scaling(
-                data_dict['gt_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE'], return_scale=True
-            )
-
-        data_dict['gt_boxes'] = gt_boxes
-        data_dict['points'] = points
+        tfs, noise_scale = augmentor_utils.global_scaling(scale_range=config['WORLD_SCALE_RANGE'], return_scale=True)
+        common_utils.apply_data_transform(data_dict, tfs)
         data_dict['noise_scale'] = noise_scale
         return data_dict
 
@@ -147,15 +118,12 @@ def random_world_translation(self, data_dict=None, config=None):
             np.random.normal(0, noise_translate_std[2], 1),
         ], dtype=np.float32).T
 
-        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
-        points[:, :3] += noise_translate
-        gt_boxes[:, :3] += noise_translate
-                
-        if 'roi_boxes' in data_dict.keys():
-            data_dict['roi_boxes'][:, :, :3] += noise_translate
+        def translate_locationlike(locations):
+            locations[..., :3] += noise_translate
+            return locations
         
-        data_dict['gt_boxes'] = gt_boxes
-        data_dict['points'] = points
+        tfs = dict(point=translate_locationlike, box=translate_locationlike)
+        common_utils.apply_data_transform(data_dict, tfs)
         data_dict['noise_translate'] = noise_translate
         return data_dict
 
diff --git a/pcdet/datasets/processor/data_processor.py b/pcdet/datasets/processor/data_processor.py
index 4f72ab532..0e5660048 100644
--- a/pcdet/datasets/processor/data_processor.py
+++ b/pcdet/datasets/processor/data_processor.py
@@ -82,7 +82,8 @@ def mask_points_and_boxes_outside_range(self, data_dict=None, config=None):
 
         if data_dict.get('points', None) is not None:
             mask = common_utils.mask_points_by_range(data_dict['points'], self.point_cloud_range)
-            data_dict['points'] = data_dict['points'][mask]
+            tfs = dict(point=lambda x: x[mask])
+            common_utils.apply_data_transform(data_dict, tfs)
 
         if data_dict.get('gt_boxes', None) is not None and config.REMOVE_OUTSIDE_BOXES and self.training:
             mask = box_utils.mask_boxes_outside_range_numpy(
@@ -97,10 +98,9 @@ def shuffle_points(self, data_dict=None, config=None):
             return partial(self.shuffle_points, config=config)
 
         if config.SHUFFLE_ENABLED[self.mode]:
-            points = data_dict['points']
-            shuffle_idx = np.random.permutation(points.shape[0])
-            points = points[shuffle_idx]
-            data_dict['points'] = points
+            shuffle_idx = np.random.permutation(data_dict['points'].shape[0])
+            tfs = dict(point=lambda x: x[shuffle_idx])
+            common_utils.apply_data_transform(data_dict, tfs)
 
         return data_dict
 
@@ -208,7 +208,8 @@ def sample_points(self, data_dict=None, config=None):
                 extra_choice = np.random.choice(choice, num_points - len(points), replace=False)
                 choice = np.concatenate((choice, extra_choice), axis=0)
             np.random.shuffle(choice)
-        data_dict['points'] = points[choice]
+        tfs = dict(point=lambda x: x[choice])
+        common_utils.apply_data_transform(data_dict, tfs)
         return data_dict
 
     def calculate_grid_size(self, data_dict=None, config=None):
diff --git a/pcdet/utils/common_utils.py b/pcdet/utils/common_utils.py
index 8864ceb46..cb573df4c 100644
--- a/pcdet/utils/common_utils.py
+++ b/pcdet/utils/common_utils.py
@@ -32,6 +32,18 @@ def drop_info_with_name(info, name):
     return ret_info
 
 
+def apply_data_transform(data_dict, transforms):
+    assert set(transforms.keys()).issubset({'point', 'box'})
+    data_keys = {
+        'point': ['points'],
+        'box': ['gt_boxes', 'roi_boxes']
+    }
+    for tf_type, tf in transforms.items():
+        for data_key in data_keys[tf_type]:
+            if data_key in data_dict:
+                data_dict[data_key] = tf(data_dict[data_key])
+
+
 def rotate_points_along_z(points, angle):
     """
     Args:

From dac816112166c35b1576cb961d4aa23aadc76f09 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Sat, 24 Jun 2023 00:12:18 +0200
Subject: [PATCH 8/9] save lidar extrinsics to dataset info pickle files

---
 pcdet/datasets/waymo/waymo_dataset.py |  2 +-
 pcdet/datasets/waymo/waymo_utils.py   | 12 ++++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/pcdet/datasets/waymo/waymo_dataset.py b/pcdet/datasets/waymo/waymo_dataset.py
index 44632d53f..12e22fbb8 100644
--- a/pcdet/datasets/waymo/waymo_dataset.py
+++ b/pcdet/datasets/waymo/waymo_dataset.py
@@ -782,7 +782,7 @@ def create_waymo_gt_database(
     parser = argparse.ArgumentParser(description='arg parser')
     parser.add_argument('--cfg_file', type=str, default=None, help='specify the config of dataset')
     parser.add_argument('--func', type=str, default='create_waymo_infos', help='')
-    parser.add_argument('--processed_data_tag', type=str, default='waymo_processed_data_v0_5_0', help='')
+    parser.add_argument('--processed_data_tag', type=str, default='waymo_processed_data_v0_6_0', help='')
     parser.add_argument('--update_info_only', action='store_true', default=False, help='')
     parser.add_argument('--use_parallel', action='store_true', default=False, help='')
     parser.add_argument('--wo_crop_gt_with_tail', action='store_true', default=False, help='')
diff --git a/pcdet/datasets/waymo/waymo_utils.py b/pcdet/datasets/waymo/waymo_utils.py
index 9bdcfe7b2..4dcb620d0 100644
--- a/pcdet/datasets/waymo/waymo_utils.py
+++ b/pcdet/datasets/waymo/waymo_utils.py
@@ -92,6 +92,7 @@ def convert_range_image_to_point_cloud(frame, range_images, camera_projections,
     points_NLZ = []
     points_intensity = []
     points_elongation = []
+    extrinsics = []
 
     frame_pose = tf.convert_to_tensor(np.reshape(np.array(frame.pose.transform), [4, 4]))
     # [H, W, 6]
@@ -162,8 +163,9 @@ def convert_range_image_to_point_cloud(frame, range_images, camera_projections,
         points_NLZ.append(np.concatenate(points_NLZ_single, axis=0))
         points_intensity.append(np.concatenate(points_intensity_single, axis=0))
         points_elongation.append(np.concatenate(points_elongation_single, axis=0))
+        extrinsics.append(extrinsic)
 
-    return points, cp_points, points_NLZ, points_intensity, points_elongation
+    return points, cp_points, points_NLZ, points_intensity, points_elongation, extrinsics
 
 
 def save_lidar_points(frame, cur_save_path, use_two_returns=True):
@@ -174,7 +176,7 @@ def save_lidar_points(frame, cur_save_path, use_two_returns=True):
         assert len(ret_outputs) == 3
         range_images, camera_projections, range_image_top_pose = ret_outputs
 
-    points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation = convert_range_image_to_point_cloud(
+    points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation, extrinsics = convert_range_image_to_point_cloud(
         frame, range_images, camera_projections, range_image_top_pose, ri_index=(0, 1) if use_two_returns else (0,)
     )
 
@@ -191,7 +193,7 @@ def save_lidar_points(frame, cur_save_path, use_two_returns=True):
 
     np.save(cur_save_path, save_points)
     # print('saving to ', cur_save_path)
-    return num_points_of_each_lidar
+    return num_points_of_each_lidar, extrinsics
 
 
 def process_single_sequence(sequence_file, save_path, sampled_interval, has_label=True, use_two_returns=True, update_info_only=False):
@@ -251,11 +253,13 @@ def process_single_sequence(sequence_file, save_path, sampled_interval, has_labe
         if update_info_only and sequence_infos_old is not None:
             assert info['frame_id'] == sequence_infos_old[cnt]['frame_id']
             num_points_of_each_lidar = sequence_infos_old[cnt]['num_points_of_each_lidar']
+            extrinsics = sequence_infos_old[cnt]['extrinsics']
         else:
-            num_points_of_each_lidar = save_lidar_points(
+            num_points_of_each_lidar, extrinsics = save_lidar_points(
                 frame, cur_save_dir / ('%04d.npy' % cnt), use_two_returns=use_two_returns
             )
         info['num_points_of_each_lidar'] = num_points_of_each_lidar
+        info['extrinsics'] = extrinsics
 
         sequence_infos.append(info)
 

From 90bc1937484af9a758ef848241a8b760add614f9 Mon Sep 17 00:00:00 2001
From: David Josef Emmerichs <david_josef.emmerichs@daimler.com>
Date: Sat, 24 Jun 2023 00:21:35 +0200
Subject: [PATCH 9/9] load pointwise laser origins

---
 pcdet/datasets/dataset.py             |  2 +-
 pcdet/datasets/waymo/waymo_dataset.py | 67 ++++++++++++++++++++-------
 pcdet/utils/common_utils.py           |  2 +-
 3 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/pcdet/datasets/dataset.py b/pcdet/datasets/dataset.py
index c1a7f6b03..0becd65f2 100644
--- a/pcdet/datasets/dataset.py
+++ b/pcdet/datasets/dataset.py
@@ -234,7 +234,7 @@ def collate_batch(batch_list, _unused=False):
                         batch_size_ratio = len(val[0])
                         val = [i for item in val for i in item]
                     ret[key] = np.concatenate(val, axis=0)
-                elif key in ['points', 'voxel_coords']:
+                elif key in ['points', 'origins', 'voxel_coords']:
                     coors = []
                     if isinstance(val[0], list):
                         val =  [i for item in val for i in item]
diff --git a/pcdet/datasets/waymo/waymo_dataset.py b/pcdet/datasets/waymo/waymo_dataset.py
index 12e22fbb8..1554edcb8 100644
--- a/pcdet/datasets/waymo/waymo_dataset.py
+++ b/pcdet/datasets/waymo/waymo_dataset.py
@@ -127,7 +127,7 @@ def load_data_to_shared_memory(self):
             if os.path.exists(f"/dev/shm/{sa_key}"):
                 continue
 
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, _origins = self.get_lidar(sequence_name, sample_idx)
             common_utils.sa_create(f"shm://{sa_key}", points)
 
         dist.barrier()
@@ -193,19 +193,34 @@ def get_infos(self, raw_data_path, save_path, num_workers=multiprocessing.cpu_co
         all_sequences_infos = [item for infos in sequence_infos for item in infos]
         return all_sequences_infos
 
-    def get_lidar(self, sequence_name, sample_idx):
+    def get_lidar(self, sequence_name, sample_idx, load_origins=False):
         lidar_file = self.data_path / sequence_name / ('%04d.npy' % sample_idx)
         point_features = np.load(lidar_file)  # (N, 7): [x, y, z, intensity, elongation, NLZ_flag]
 
         points_all, NLZ_flag = point_features[:, 0:5], point_features[:, 5]
         if not self.dataset_cfg.get('DISABLE_NLZ_FLAG_ON_POINTS', False):
             points_all = points_all[NLZ_flag == -1]
+
         if self.dataset_cfg.get('POINTS_TANH_DIM', None) is None:
             points_all[:, 3] = np.tanh(points_all[:, 3])
         else:
             for dim_idx in self.dataset_cfg.POINTS_TANH_DIM:
                 points_all[:, dim_idx] = np.tanh(points_all[:, dim_idx])
-        return points_all
+
+        if load_origins:
+            sample_info = self.seq_name_to_infos[sequence_name][sample_idx]
+            if 'extrinsics' not in sample_info:
+                raise ValueError('extrinsics not saved to database, use db version >= v0_6_0')
+            origins = [extr[:3, 3] for extr in sample_info['extrinsics']]
+            laser_counts = sample_info['num_points_of_each_lidar']
+            assert sum(laser_counts) == points_all.shape[0], (laser_counts, points_all.shape)
+            assert len(origins) == len(laser_counts), (origins, laser_counts)
+            origins = np.concatenate([np.tile(extr[None, :], (c, 1)) for c, extr in zip(laser_counts, origins)], axis=0)
+            assert origins.shape == points_all[:, :3].shape, (origins.shape, points_all.shape)
+        else:
+            origins = None
+
+        return points_all, origins
 
     @staticmethod
     def transform_prebox_to_current(pred_boxes3d, pose_pre, pose_cur):
@@ -247,7 +262,7 @@ def reorder_rois_for_refining(pred_bboxes):
             ordered_bboxes[bs_idx, :len(pred_bboxes[bs_idx])] = pred_bboxes[bs_idx]
         return ordered_bboxes
 
-    def get_sequence_data(self, info, points, sequence_name, sample_idx, sequence_cfg, load_pred_boxes=False):
+    def get_sequence_data(self, info, points, origins, sequence_name, sample_idx, sequence_cfg, load_pred_boxes=False):
         """
         Args:
             info:
@@ -260,7 +275,7 @@ def get_sequence_data(self, info, points, sequence_name, sample_idx, sequence_cf
 
         def remove_ego_points(points, center_radius=1.0):
             mask = ~((np.abs(points[:, 0]) < center_radius) & (np.abs(points[:, 1]) < center_radius))
-            return points[mask]
+            return points[mask], mask
 
         def load_pred_boxes_from_dict(sequence_name, sample_idx):
             """
@@ -272,6 +287,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             load_boxes[:, 7:9] = -0.1 * load_boxes[:, 7:9]  # transfer speed to negtive motion from t to t-1
             return load_boxes
 
+        load_origins = origins is not None
         pose_cur = info['pose'].reshape((4, 4))
         num_pts_cur = points.shape[0]
         sample_idx_pre_list = np.clip(sample_idx + np.arange(sequence_cfg.SAMPLE_OFFSET[0], sequence_cfg.SAMPLE_OFFSET[1]), 0, 0x7FFFFFFF)
@@ -285,6 +301,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             points = np.hstack([points, np.zeros((points.shape[0], 1)).astype(points.dtype)])
         points_pre_all = []
         num_points_pre = []
+        origins_pre_all = []
 
         pose_all = [pose_cur]
         pred_boxes_all = []
@@ -296,7 +313,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
 
         for idx, sample_idx_pre in enumerate(sample_idx_pre_list):
 
-            points_pre = self.get_lidar(sequence_name, sample_idx_pre)
+            points_pre, origins_pre = self.get_lidar(sequence_name, sample_idx_pre, load_origins=load_origins)
             pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4))
             expand_points_pre = np.concatenate([points_pre[:, :3], np.ones((points_pre.shape[0], 1))], axis=-1)
             points_pre_global = np.dot(expand_points_pre, pose_pre.T)[:, :3]
@@ -310,11 +327,19 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
             else:
                 # add timestamp
                 points_pre = np.hstack([points_pre, 0.1 * (sample_idx - sample_idx_pre) * np.ones((points_pre.shape[0], 1)).astype(points_pre.dtype)])  # one frame 0.1s
-            points_pre = remove_ego_points(points_pre, 1.0)
+            points_pre, ego_mask = remove_ego_points(points_pre, 1.0)
             points_pre_all.append(points_pre)
             num_points_pre.append(points_pre.shape[0])
             pose_all.append(pose_pre)
 
+            if load_origins:
+                expand_origins_pre = np.concatenate([origins_pre[:, :3], np.ones((origins_pre.shape[0], 1))], axis=-1)
+                origins_pre_global = np.dot(expand_origins_pre, pose_pre.T)[:, :3]
+                expand_origins_pre_global = np.concatenate([origins_pre_global, np.ones((origins_pre_global.shape[0], 1))], axis=-1)
+                origins_pre = np.dot(expand_origins_pre_global, np.linalg.inv(pose_cur.T))[:, :3]
+                origins_pre = origins_pre[ego_mask]
+                origins_pre_all.append(origins_pre)
+
             if load_pred_boxes:
                 pose_pre = sequence_info[sample_idx_pre]['pose'].reshape((4, 4))
                 pred_boxes = load_pred_boxes_from_dict(sequence_name, sample_idx_pre)
@@ -325,6 +350,11 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
         num_points_all = np.array([num_pts_cur] + num_points_pre).astype(np.int32)
         poses = np.concatenate(pose_all, axis=0).astype(np.float32)
 
+        if load_origins:
+            origins = np.concatenate([origins] + origins_pre_all, axis=0).astype(np.float32)
+        else:
+            origins = None
+
         if load_pred_boxes:
             temp_pred_boxes = self.reorder_rois_for_refining(pred_boxes_all)
             pred_boxes = temp_pred_boxes[:, :, 0:9]
@@ -333,7 +363,7 @@ def load_pred_boxes_from_dict(sequence_name, sample_idx):
         else:
             pred_boxes = pred_scores = pred_labels = None
 
-        return points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels
+        return points, origins, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels
 
     def __len__(self):
         if self._merge_all_iters_to_one_epoch:
@@ -352,15 +382,15 @@ def __getitem__(self, index):
         input_dict = {
             'sample_idx': sample_idx
         }
-        if self.use_shared_memory and index < self.shared_memory_file_limit:
+        if self.use_shared_memory and index < self.shared_memory_file_limit and not self.dataset_cfg.get('USE_ORIGINS', False):
             sa_key = f'{sequence_name}___{sample_idx}'
             points = SharedArray.attach(f"shm://{sa_key}").copy()
         else:
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, origins = self.get_lidar(sequence_name, sample_idx, load_origins=self.dataset_cfg.get('USE_ORIGINS', False))
 
         if self.dataset_cfg.get('SEQUENCE_CONFIG', None) is not None and self.dataset_cfg.SEQUENCE_CONFIG.ENABLED:
-            points, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels = self.get_sequence_data(
-                info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG,
+            points, origins, num_points_all, sample_idx_pre_list, poses, pred_boxes, pred_scores, pred_labels = self.get_sequence_data(
+                info, points, origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG,
                 load_pred_boxes=self.dataset_cfg.get('USE_PREDBOX', False)
             )
             input_dict['poses'] = poses
@@ -373,6 +403,7 @@ def __getitem__(self, index):
 
         input_dict.update({
             'points': points,
+            'origins': origins,
             'frame_id': info['frame_id'],
         })
 
@@ -491,11 +522,11 @@ def create_groundtruth_database(self, info_path, save_path, used_classes=None, s
             pc_info = info['point_cloud']
             sequence_name = pc_info['lidar_sequence']
             sample_idx = pc_info['sample_idx']
-            points = self.get_lidar(sequence_name, sample_idx)
+            points, _origins = self.get_lidar(sequence_name, sample_idx)
 
             if use_sequence_data:
-                points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
-                    info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
+                points, _origins, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
+                    info, points, _origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
                 )
 
             annos = info['annos']
@@ -569,11 +600,11 @@ def create_gt_database_of_single_scene(self, info_with_idx, database_save_path=N
         pc_info = info['point_cloud']
         sequence_name = pc_info['lidar_sequence']
         sample_idx = pc_info['sample_idx']
-        points = self.get_lidar(sequence_name, sample_idx)
+        points, _origins = self.get_lidar(sequence_name, sample_idx)
 
         if use_sequence_data:
-            points, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
-                info, points, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
+            points, _origins, num_points_all, sample_idx_pre_list, _, _, _, _ = self.get_sequence_data(
+                info, points, _origins, sequence_name, sample_idx, self.dataset_cfg.SEQUENCE_CONFIG
             )
 
         annos = info['annos']
diff --git a/pcdet/utils/common_utils.py b/pcdet/utils/common_utils.py
index cb573df4c..caa57d94d 100644
--- a/pcdet/utils/common_utils.py
+++ b/pcdet/utils/common_utils.py
@@ -35,7 +35,7 @@ def drop_info_with_name(info, name):
 def apply_data_transform(data_dict, transforms):
     assert set(transforms.keys()).issubset({'point', 'box'})
     data_keys = {
-        'point': ['points'],
+        'point': ['points', 'origins'],
         'box': ['gt_boxes', 'roi_boxes']
     }
     for tf_type, tf in transforms.items():