More preprocess changes + video batching

Author: thetianshuhuang

dart/preprocess/dataset.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from beartype.typing import NamedTuple
-from jaxtyping import Complex64, Int16
+from jaxtyping import Complex64, Int16, Float64
 
 from .radar import AWR1843Boost
 
@@ -52,7 +52,7 @@ def _get_frames(self, packets, valid):
         res = rad[self.start:self.end].reshape(-1, self.frame_size)[valid]
         return res
 
-    def _get_times(self, packets, valid):
+    def _get_times(self, packets, valid) -> Float64[np.ndarray, "frames"]:
         """Get timestamps for each frame.
 
         Timestamps are denoted by the first packet corresponding to data
@@ -75,7 +75,12 @@ def _to_iq(
         iq[:, 1::2] = frames[:, 1::4] + 1j * frames[:, 3::4]
         return iq.reshape((-1, *radar.frame_shape))
 
-    def as_frames(self, radar, packets):
+    def as_frames(
+        self, radar: AWR1843Boost, packets
+    ) -> tuple[
+        Complex64[np.ndarray, "frames tx rx chirp"],
+        Float64[np.ndarray, "frames"]
+    ]:
         """Convert packets to frames."""
         valid = self._get_valid(packets)
         data = self._get_frames(packets, valid)

dart/preprocess/radar.py

@@ -13,8 +13,8 @@
 
 
 def to_float16(
-    arr: Float32[Array, "..."], max_normal: float = 65504.0
-) -> Float16[Array, "..."]:
+    arr: Float32[types.ArrayLike, "..."], max_normal: float = 65504.0
+) -> Float16[types.ArrayLike, "..."]:
     """Convert array to float16.
 
     The largest observed value is scaled to the largest normal normal number
@@ -79,14 +79,18 @@ def estimate_speed(
         """Estimate speed for this frame using doppler spectrum."""
         nd_half = int(self.framelen / 2)
 
-        # Spectrum across doppler bins; folded in half
-        _spectrum = np.sum(rda, axis=(1, 3))
-        spectrum = _spectrum[:, :nd_half] + _spectrum[:, nd_half:][:, ::-1]
+        percentile = 99.75
+        min_count = 10
 
-        threshold = jnp.min(spectrum, axis=1).reshape(-1, 1)
-        speed_nd = jnp.argmax(
-            (jnp.diff(spectrum, axis=1) > threshold)
-            & (spectrum[:, :-1] > (threshold * 2)), axis=1)
+        # Threshold for each frame, antenna (over=1,2)
+        threshold = jnp.percentile(
+            rda, percentile, axis=(1, 2))[:, None, None, ...]
+        # Count across range, antenna (over=2,3)
+        valid = jnp.sum(rda > threshold, axis=(1, 3)) > min_count
+
+        left = jnp.argmax(valid, axis=1)
+        right = jnp.argmax(valid[:, ::-1], axis=1)
+        speed_nd = jnp.minimum(left, right)
 
         return (nd_half - speed_nd) / nd_half * self.dmax / 2
 
@@ -126,8 +130,8 @@ def range_doppler_azimuth(
         return rda[:, :self.max_range], self.estimate_speed(rda)
 
     def remove_artifact(
-        self, images: Float32[types.ArrayLike, "frame range doppler antenna"],
-    ) -> Float32[types.ArrayLike, "frame range doppler antenna"]:
+        self, images: Float32[Array, "frame range doppler antenna"],
+    ) -> Float32[Array, "frame range doppler antenna"]:
         """Remove zero-doppler artifact from images.
 
         Collected range-doppler radar data will have an artifact at close
@@ -138,7 +142,7 @@ def remove_artifact(
         zero = int(images.shape[2] / 2)
         artifact = jnp.percentile(
             images[..., zero - 1:zero + 2, :], self.artifact_threshold, axis=0)
-        removed = np.maximum(
+        removed = jnp.maximum(
             images[..., zero - 1:zero + 2, :]
             - artifact.reshape(1, *artifact.shape), 0)
         return images.at[..., zero - 1:zero + 2, :].set(removed)
@@ -197,6 +201,7 @@ def process_data(
             chirps_flat, window_shape=self.framelen, axis=0)[::stride]
 
         process_func = jax.jit(partial(self.range_doppler_azimuth))
+        # process_func = partial(self.range_doppler_azimuth)
         res, speed = [], []
         for _ in range(int(np.ceil(frames.shape[0] / batch_size))):
             r, s = process_func(jnp.array(frames[:batch_size]))

dart/preprocess/trajectory.py

@@ -10,6 +10,8 @@
 from beartype.typing import NamedTuple
 from jaxtyping import Float32, Float64, Bool
 
+from dart import types
+
 
 class Trajectory(NamedTuple):
     """Sensor trajectory.
@@ -36,10 +38,8 @@ def from_csv(cls, path: str) -> "Trajectory":
         """
         df = pd.read_csv(os.path.join(path, "trajectory.csv"))
 
-        t_slam = np.array(df["field.header.stamp"]) / 1e9
-
-        # tmp
-        t_slam = t_slam + 0.5
+        # Manual 0.5s offset, likely due to buffering the DCA1000
+        t_slam = np.array(df["field.header.stamp"]) / 1e9 + 0.5
 
         xyz = np.array(
             [df["field.transform.translation." + char] for char in "xyz"])
@@ -51,14 +51,18 @@ def from_csv(cls, path: str) -> "Trajectory":
 
         return cls(spline=spline, slerp=slerp)
 
-    def valid_mask(self, t: Float64[np.ndarray, "N"]) -> Bool[np.ndarray, "N"]:
+    def valid_mask(
+        self, t: Float64[types.ArrayLike, "N"], window: float = 0.1
+    ) -> Bool[types.ArrayLike, "N"]:
         """Get mask of valid timestamps (within the trajectory timestamps)."""
-        return (t >= self.spline.x[0]) & (t <= self.spline.x[-1])
+        return (
+            (t - window >= self.spline.x[0])
+            & (t + window <= self.spline.x[-1]))
 
     def interpolate(
-        self, t: Float64[np.ndarray, "N"], window: float = 0.1,
+        self, t: Float64[types.ArrayLike, "N"], window: float = 0.1,
         samples: int = 25
-    ) -> dict[str, Float32[np.ndarray, "N ..."]]:
+    ) -> dict[str, Float32[types.ArrayLike, "N ..."]]:
         """Calculate poses, averaging along a window.
 
         Parameters
@@ -72,13 +76,13 @@ def interpolate(
         Dictionary with pos, vel, and rot entries.
         """
         window_offsets = np.linspace(-window / 2, window / 2, samples)
-        samples = t[..., None] + window_offsets[None, ...]
+        tt = t[..., None] + window_offsets[None, ...]
 
         # Rotation unfortunately does not allow vectorization at this time.
-        rot = Rotation.concatenate([self.slerp(row).mean() for row in samples])
+        rot = Rotation.concatenate([self.slerp(row).mean() for row in tt])
 
         return {
-            "pos": np.mean(self.spline(samples), axis=1),
-            "vel": np.mean(self.spline.derivative()(samples), axis=1),
+            "pos": np.mean(self.spline(tt), axis=1),
+            "vel": np.mean(self.spline.derivative()(tt), axis=1),
             "rot": rot.as_matrix()
         }

dart/result.py

@@ -5,11 +5,13 @@
 import numpy as np
 import h5py
 import matplotlib as mpl
+from jax import numpy as jnp
 
 from beartype.typing import Optional, Any
 from jaxtyping import Integer, Array, Float, UInt8
 
 from .dart import DART
+from dart.jaxcolors import colormap
 from .dataset import load_arrays, trajectory
 from . import types
 
@@ -39,6 +41,8 @@ def __init__(self, path: str) -> None:
         with open(_meta) as f:
             self.metadata = json.load(f)
 
+        self.DATASET = self.metadata["dataset"]["path"]
+
     def dart(self) -> DART:
         """Get DART object."""
         return DART.from_config(**self.metadata)
@@ -64,6 +68,10 @@ def load(self, path: str, keys: Optional[list[str]] = None) -> Any:
         """Load file inside this result."""
         return load_arrays(os.path.join(self.path, path), keys=keys)
 
+    def open(self, path: str) -> Any:
+        """Load h5py file in this scope."""
+        return h5py.File(os.path.join(self.path, path))
+
     @staticmethod
     def colorize_map(
         arr: Float[types.ArrayLike, "..."], sigma: bool = True,
@@ -99,3 +107,35 @@ def colorize_map(
             arr = np.exp(arr)  # type: ignore
 
         return (mpl.colormaps['viridis'](arr)[..., :3] * 255).astype(np.uint8)
+
+    @staticmethod
+    def colorize_radar(
+        cmap: Float[types.ArrayLike, "..."],
+        rad: Float[types.ArrayLike, "..."],
+        clip: tuple[float, float] = (5.0, 99.9)
+    ) -> UInt8[types.ArrayLike, "... 3"]:
+        """Colorize a radar intensity map in a jax-friendly way.
+
+        Parameters
+        ----------
+        cmap: color map to apply.
+        rad: input array. If range-doppler, colorize directly; if
+            range-doppler-azimuth, tiles into 2 columns x 4 rows.
+        clip: percentile clipping range.
+        """
+        def _tile(arr):
+            unpack = [arr[:, :, :, i] for i in range(arr.shape[3])]
+            left = jnp.concatenate(unpack[:4], axis=1)
+            right = jnp.concatenate(unpack[4:], axis=1)
+            return jnp.concatenate([left, right], axis=2)
+
+        p5, p95 = jnp.nanpercentile(rad, jnp.array(clip))
+        rad = (rad - p5) / (p95 - p5)
+        colors = (colormap(cmap, rad) * 255).astype(jnp.uint8)
+        if len(rad.shape) == 4:
+            if rad.shape[-1] > 1:
+                return _tile(colors)
+            else:
+                return colors[..., 0, :]
+        else:
+            return colors

notebooks/map.ipynb

@@ -5,6 +5,7 @@
 import json
 from tqdm import tqdm
 import numpy as np
+from scipy.ndimage import gaussian_filter1d
 
 from dart.preprocess import AWR1843Boost, AWR1843BoostDataset, Trajectory
 
@@ -27,7 +28,7 @@ def _process_batch(radar: AWR1843Boost, file_batch, traj: Trajectory):
     range_doppler, speed_est = radar.process_data(chirps)
     t_image = radar.process_timestamps(t_chirp)
 
-    t_valid = traj.valid_mask(t_image)
+    t_valid = traj.valid_mask(t_image, window=radar.frame_time * 0.5)
     pose = traj.interpolate(t_image[t_valid], window=radar.frame_time * 0.5)
     pose["t"] = t_image[t_valid]
     pose["speed"] = speed_est[t_valid]
@@ -37,7 +38,7 @@ def _process_batch(radar: AWR1843Boost, file_batch, traj: Trajectory):
 
 def _process(
         path: str, radar: AWR1843Boost, batch_size: int = 1000000,
-        overwrite: bool = False):
+        overwrite: bool = False, sigma: float = 2.0):
 
     traj = Trajectory.from_csv(path)
     packetfile = h5py.File(os.path.join(path, "radarpackets.h5"), 'r')
@@ -53,22 +54,25 @@ def _process(
         "rad", (1, *rs), dtype='f2', chunks=(1, *rs), maxshape=(None, *rs))
 
     total_size = 0
-    poses = []
+    _poses = []
     for _ in tqdm(range(int(np.ceil(packet_dataset.shape[0] / batch_size)))):
         rda, pose = _process_batch(radar, packet_dataset[:batch_size], traj)
-        poses.append(pose)
+        _poses.append(pose)
 
         total_size += rda.shape[0]
         range_doppler_azimuth.resize((total_size, *radar.image_shape))
         range_doppler_azimuth[-rda.shape[0]:] = rda
 
         packet_dataset = packet_dataset[batch_size:]
 
-    poses_cat = {}
-    for k in poses[0]:
-        poses_cat[k] = np.concatenate([p[k] for p in poses])
+    poses = {}
+    for k in _poses[0]:
+        poses[k] = np.concatenate([p[k] for p in _poses])
 
-    for k, v in poses_cat.items():
+    if sigma > 0:
+        poses['vel'] = gaussian_filter1d(poses['vel'], sigma=sigma, axis=0)
+
+    for k, v in poses.items():
         outfile.create_dataset(k, data=v)
     packetfile.close()
     outfile.close()