fix(visualize_pose): vis fixes

.github/workflows/main.yml

@@ -52,41 +52,3 @@ jobs:
         uses: codecov/codecov-action@v3
         # with:
         #   fail_ci_if_error: true
-
-  tests_mac:
-    needs: linter
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: [3.9]
-        os: [macos-latest]
-    runs-on: ${{ matrix.os }}
-    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install project
-        run: make install
-      - name: Run tests
-        run: make test
-
-  tests_win:
-    needs: linter
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: [3.9]
-        os: [windows-latest]
-    runs-on: ${{ matrix.os }}
-    steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
-        with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install Pip
-        run: pip install --user --upgrade pip
-      - name: Install project
-        run: pip install -e .[test]
-      - name: run tests
-        run: pytest -s -vvvv -l --tb=long tests

README.md

@@ -31,19 +31,24 @@ $ drivellava
 
 Read the [CONTRIBUTING.md](CONTRIBUTING.md) file.
 
-## Dataset setup
+## Running the scripts
 
 ```python3
 python3 -m drivellava.scripts.generate_commavq_images
 ```
 
+```python3
+python3 -m drivellava.scripts.visualize_pose
+```
+
 ## TODO
 
 - Training script
   - Select what layers to train
   - Measure memory requirements
 - Dataset
   - Generate images from CommaVQ
+  - Denoise the trajectory
   - Quantize the trajectory
   - Visualize the trajectory on the image
   - Generate JSON dataset

drivellava/constants.py

@@ -24,6 +24,7 @@ def __getitem__(self, index):
 # List of all the videos
 ENCODED_VIDEOS_ALL = glob(os.path.join(COMMAVQ_DIR, "*", "*.npy"))
 ENCODED_VIDEOS_ALL = [x for x in ENCODED_VIDEOS_ALL if os.path.isfile(x)]
+# ENCODED_VIDEOS_ALL = sorted(ENCODED_VIDEOS_ALL)
 
 ENCODED_POSE = glob(os.path.join(COMMAVQ_DIR, "*", "*.npy"))
 ENCODED_POSE = [x for x in ENCODED_POSE if os.path.isfile(x)]

drivellava/scripts/visualize_pose.py

@@ -12,10 +12,44 @@
 from drivellava.utils import (
     plot_bev_trajectory,
     plot_steering_traj,
-    smoothen_traj,
+    remove_noise,
 )
 
 
+def get_rotation_matrix(pose):
+    """
+    origin: (x, y, z, roll, pitch, yaw)
+    """
+    # Construct 3D rotation matrix
+    origin = pose.copy()
+
+    rot_x = np.array(
+        [
+            [1, 0, 0],
+            [0, np.cos(origin[3]), -np.sin(origin[3])],
+            [0, np.sin(origin[3]), np.cos(origin[3])],
+        ]
+    )
+    rot_y = np.array(
+        [
+            [np.cos(origin[4]), 0, np.sin(origin[4])],
+            [0, 1, 0],
+            [-np.sin(origin[4]), 0, np.cos(origin[4])],
+        ]
+    )
+    rot_z = np.array(
+        [
+            [np.cos(origin[5]), -np.sin(origin[5]), 0],
+            [np.sin(origin[5]), np.cos(origin[5]), 0],
+            [0, 0, 1],
+        ]
+    )
+
+    rot = np.dot(rot_z, np.dot(rot_y, rot_x))
+
+    return rot
+
+
 def get_local_pose(
     pose: np.ndarray,
     frame_index: int,
@@ -36,67 +70,48 @@ def get_local_pose(
     sub_pose[0, :] = 0
 
     # Each element is (x, y, z, roll, pitch, yaw)
-
-    # # Construct 3D rotation matrix
-    # rot_x = np.array(
-    #     [
-    #         [1, 0, 0],
-    #         [0, np.cos(origin[3]), -np.sin(origin[3])],
-    #         [0, np.sin(origin[3]), np.cos(origin[3])],
-    #     ]
-    # )
-    # rot_y = np.array(
-    #     [
-    #         [np.cos(origin[4]), 0, np.sin(origin[4])],
-    #         [0, 1, 0],
-    #         [-np.sin(origin[4]), 0, np.cos(origin[4])],
-    #     ]
-    # )
-    # rot_z = np.array(
-    #     [
-    #         [np.cos(origin[5]), -np.sin(origin[5]), 0],
-    #         [np.sin(origin[5]), np.cos(origin[5]), 0],
-    #         [0, 0, 1],
-    #     ]
-    # )
-
-    # rot = np.dot(rot_z, np.dot(rot_y, rot_x))
-
-    # Each pose element is the relative transformation from the previous pose
-    # to the current pose. We need to accumulate the transformations to get the
-    # absolute pose
-    # for i in range(1, sub_pose.shape[0]):
-    #     sub_pose[i, :3] += sub_pose[i - 1, :3]
+    for i in range(1, sub_pose.shape[0]):
+        rot = get_rotation_matrix(sub_pose[i])
+        sub_pose[i, :3] += np.dot(sub_pose[:, :3], rot)[i - 1, :3]
 
     sub_pose[1:, :3] += sub_pose[:-1, :3]
 
-    # Rotate and translate the pose with respect to the origin
-    # sub_pose = sub_pose - origin
-    # sub_pose[:, :3] = np.dot(sub_pose[:, :3], rot)
-
     return sub_pose
 
 
 def main():
 
-    NUM_FRAMES = 20 * 5
+    NUM_FRAMES = 20 * 1
 
     for encoded_video_path in tqdm(
         DECODED_IMGS_ALL_AVAILABLE.keys(), desc="npy files"
     ):
         decoded_imgs_list = DECODED_IMGS_ALL_AVAILABLE[encoded_video_path]
-        pose_path = encoded_video_path.replace("data_", "pose_").replace(
+        pose_path = encoded_video_path.replace("data_", "pose_data_").replace(
             "val", "pose_val"
         )
         print(pose_path)
 
-        assert os.path.exists(pose_path)
+        assert os.path.exists(pose_path), pose_path
 
         pose = np.load(pose_path)
 
+        # Inches to meters
+        pose[:, :3] *= 0.0254
+        pose[:, 1] = 0.0
+        # pose[:, 2] *= 0.01
+
         # pose[1:, :3] += pose[:-1, :3]
 
         # pose[:, :3] = smoothen_traj(pose[:, :3])
+        # pose = smoothen_traj(pose, window_size=100)
+        # pose = smoothen_traj(pose, window_size=20*3)
+        pose = remove_noise(
+            pose,
+            window_length=21 * 2 -1,
+            # window_length=9,
+            polyorder=1,
+        )
 
         pose[np.isnan(pose)] = 0
         pose[np.isinf(pose)] = 0
@@ -127,7 +142,24 @@ def main():
             # Swap axes
             trajectory = trajectory[:, [2, 1, 0]]
 
-            trajectory = smoothen_traj(trajectory)
+            print(
+                "trajectory[0]",
+                (np.min(trajectory[:, 0]), np.max(trajectory[:, 0])),
+            )
+            print(
+                "trajectory[1]",
+                (np.min(trajectory[:, 1]), np.max(trajectory[:, 1])),
+            )
+            print(
+                "trajectory[2]",
+                (np.min(trajectory[:, 2]), np.max(trajectory[:, 2])),
+            )
+            dx = trajectory[1:, 2] - trajectory[:-1, 2]
+            speed = dx / (1.0 / 20.0)
+            # m/s to km/h
+            speed_kmph = speed * 3.6
+            # speed mph
+            speed_mph = speed_kmph * 0.621371
 
             img = plot_steering_traj(
                 img,
@@ -136,7 +168,26 @@ def main():
 
             img_bev = plot_bev_trajectory(trajectory, img)
 
-            cv2.imshow("frame_path", cv2.resize(img, (0, 0), fx=2, fy=2))
+            # Write speed on img
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            bottomLeftCornerOfText = (10, 50)
+            fontScale = 0.5
+            fontColor = (255, 255, 255)
+            lineType = 2
+
+            img = cv2.resize(img, (0, 0), fx=2, fy=2)
+
+            cv2.putText(
+                img,
+                f"Speed: {speed_mph.mean():.2f} mph",
+                bottomLeftCornerOfText,
+                font,
+                fontScale,
+                fontColor,
+                lineType,
+            )
+
+            cv2.imshow("frame_path", img)
 
             cv2.imshow(
                 "frame_path_bev", cv2.resize(img_bev, (0, 0), fx=2, fy=2)