Improve obstacle tracking

uavf_2025/perception/dev/__init__.py

@@ -0,0 +1 @@
+from .array_logger import ArrayLogger as ArrayLogger

uavf_2025/perception/dev/array_logger.py

@@ -0,0 +1,66 @@
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from time import time
+from warnings import warn
+
+import numpy as np
+import torch
+
+
+class ArrayLogger:
+    """
+    Logs a numpy array as a .npy file every `period` seconds.
+    """
+
+    def __init__(self, logs_dir: Path, period: int = 1):
+        """
+        Initializes the ParticlesLogger with a directory to save logs.
+
+        Parameters
+        ----------
+        logs_dir : Path
+            Directory where the particle logs will be saved.
+        """
+        self._logs_dir = logs_dir
+        self._period = period
+
+        self._last_log_time: float = 0
+
+        self._logs_dir.mkdir(parents=True, exist_ok=True)
+
+        self._executor = ThreadPoolExecutor(max_workers=1)
+
+    def update(self, array: torch.Tensor | np.ndarray) -> None:
+        """
+        Saves the provided array to a .npy file if the specified period has passed since the last log.
+
+        The array should NOT be copied before calling this method.
+        The copy will only be made if necessary.
+
+        Parameters
+        ----------
+        array : torch.Tensor | np.ndarray
+            The array to log. If a torch tensor is provided, it will be converted to a numpy array.
+        """
+        current_time = time()
+
+        if current_time - self._last_log_time >= self._period:
+            self._last_log_time = current_time
+        else:
+            return
+
+        if isinstance(array, torch.Tensor):
+            array = array.cpu().numpy()
+        elif isinstance(array, np.ndarray):
+            array = array.copy()
+        else:
+            raise TypeError(
+                "Expected a torch.Tensor or np.ndarray, got {}".format(type(array))
+            )
+
+        filename = self._logs_dir / f"{int(current_time)}.npy"
+        if not filename.exists():
+            self._executor.submit(lambda: np.save(filename, array))
+            print(f"ArrayLogger: Saved array to {filename}")
+        else:
+            warn(f"ArrayLogger: File {filename} already exists. Not overwriting.")

uavf_2025/perception/dev/perception_cli.py

@@ -13,14 +13,17 @@
 
 
 def visualize_obstacles(drone_pose: Pose, obstacles: list[Obstacle]):
+    if len(obstacles) == 0:
+        return
+
     canvas = np.zeros((1000, 1000, 3), dtype=np.uint8)
 
     # scale image such that it covers at least 5x5 meters, and always includes every obstacle with a 1 meter buffer
     scale = max(
         5,
         max(
-            [obstacle.state[0] for obstacle in obstacles]
-            + [obstacle.state[1] for obstacle in obstacles]
+            [np.abs(obstacle.state[0]) for obstacle in obstacles]
+            + [np.abs(obstacle.state[1]) for obstacle in obstacles]
         )
         + 1,
     )

uavf_2025/perception/dev/visualizations/__init__.py

@@ -0,0 +1,2 @@
+from .particles import ParticlesVisualization as ParticlesVisualization
+from .obstacle_tracking import ObstacleTrackerVisualizer as ObstacleTrackerVisualizer

uavf_2025/perception/dev/visualizations/obstacle_tracking.py

@@ -9,7 +9,6 @@
 from perception.lib.yolo import TiledYolo
 import torch
 from tqdm import tqdm
-from matplotlib import widgets, pyplot as plt
 
 from perception.camera import DummyCamera, ImageMetadata
 from perception.obstacle_tracking import (
@@ -20,6 +19,8 @@
 from perception.types import Bbox2D, Image
 from shared.types import Pose
 
+from .particles import ParticlesVisualization
+
 
 class ObstacleTrackerVisualizer:
     def __init__(
@@ -76,49 +77,10 @@ def visualize(self) -> None:
         if self._states is None:
             raise RuntimeError("Simulation must be run before visualization")
 
-        fig = plt.figure()
-
-        ax1 = fig.add_axes((0, 0, 1, 0.8), projection="3d")  # 3D plot axis
-
-        ax2 = fig.add_axes((0.1, 0.85, 0.8, 0.1))  # slider axis
-        slider = widgets.Slider(ax2, "Time Step", 0, len(self._states) - 1, valinit=0)
-
-        if self._frames is not None:
-            assert len(self._frames) == self._iterations
-
-            ax3 = fig.add_axes((0.1, 0.7, 0.8, 0.15))  # camera frame axis
-            ax3.set_title("Camera Frame")
-
-        def update(iteration: float):
-            if self._states is None:
-                raise RuntimeError("Simulation must be run before visualization")
-
-            iteration = round(iteration)
-            if iteration < 0:
-                iteration = 0
-
-            ax1.clear()
-            ax1.set_title(f"Iteration {iteration}")
-
-            particles = self._states[iteration]
-            ax1.scatter(particles[:, 0], particles[:, 1], particles[:, 2])
-
-            drone_pose = self._drone_poses[iteration]
-            direction = drone_pose.rotation.apply([0, 0, 1])
-
-            arrow_length = torch.mean(torch.tensor(particles).norm(dim=1)).item() / 10
-            ax1.quiver(*drone_pose.position, *direction, length=arrow_length, color="r")
-
-            if self._frames is not None:
-                assert ax3, "Axis should have been created"  # type: ignore
-                ax3.clear()
-                ax3.imshow(self._frames[iteration])
-                ax3.axis("off")
-
-        slider.on_changed(update)
-        update(0)
-
-        plt.show()
+        visualizer = ParticlesVisualization(
+            self._states, self._drone_poses, self._frames
+        )
+        visualizer.visualize()
 
     def _simulate(self) -> list[np.ndarray]:
         states: list[np.ndarray] = []
@@ -129,7 +91,10 @@ def _simulate(self) -> list[np.ndarray]:
             desc="Simulating particle states",
         ):
             self._tracker.update(drone_pose)
-            states.append(self._tracker.get_particles().cpu().numpy())
+            particles = self._tracker.get_particles(True)
+            if particles is None:
+                raise RuntimeError("Tracker did not return any particles")
+            states.append(particles.cpu().numpy())
 
         return states
 
@@ -175,12 +140,17 @@ def video_tracking(source_dir: Path, iterations: int = 100):
     ObstacleTrackerVisualizer(tracker, list(drone_poses), iterations)
 
 
-def dummy_video_tracking(height: int = 1080, width: int = 1920):
+def dummy_video_tracking(
+    height: int = 1080,
+    width: int = 1920,
+    dropout: float = 0.1,  # Probability of a false negative
+):
     """
     Generates a video and corresponding bounding boxes for a square moving from
     left to right while getting smaller (i.e., equivalent to moving diagonally away),
     then runs the tracking algorithm and visualization on it.
     """
+    rng = np.random.default_rng(42)  # For reproducibility
     num_frames = 100
 
     centers = np.linspace(0.25, 0.75, num_frames)
@@ -191,16 +161,22 @@ def dummy_video_tracking(height: int = 1080, width: int = 1920):
 
     for center, size in zip(centers, sizes):
         img = np.zeros((height, width, 3), dtype=np.uint8)
-        cv2.rectangle(
-            img,
-            (int((center - size / 2) * width), int((center - size / 2) * height)),
-            (int((center + size / 2) * width), int((center + size / 2) * height)),
-            (255, 255, 255),
-            -1,
-        )
 
+        if rng.random() < dropout:
+            # Simulate a false negative by not adding a bounding box
+            detections = []
+        else:
+            cv2.rectangle(
+                img,
+                (int((center - size / 2) * width), int((center - size / 2) * height)),
+                (int((center + size / 2) * width), int((center + size / 2) * height)),
+                (255, 255, 255),
+                -1,
+            )
+            detections = [Bbox2D(center, center, size, size)]
+
+        bboxes.append(detections)
         images.append(Image(img))
-        bboxes.append([Bbox2D(center, center, size, size)])
 
     camera = DummyCamera(
         images, [_dummy_metadata_factory("") for _ in range(len(images))]
@@ -213,7 +189,9 @@ def dummy_video_tracking(height: int = 1080, width: int = 1920):
     tracker = ObstacleTracker([camera], detector, localizer)
 
     drone_poses = list(repeat(Pose.identity(), num_frames))
-    ObstacleTrackerVisualizer(tracker, drone_poses, num_frames)
+    ObstacleTrackerVisualizer(
+        tracker, drone_poses, num_frames, [image.get_array() for image in images]
+    )
 
 
 def _dummy_metadata_factory(name: str) -> ImageMetadata:
@@ -245,4 +223,5 @@ def make_and_save_metadata(img_path: Path):
 
 
 if __name__ == "__main__":
-    static_tracking()
+    # static_tracking()
+    dummy_video_tracking()

uavf_2025/perception/dev/visualizations/particles.py

@@ -0,0 +1,111 @@
+from itertools import repeat
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from matplotlib import widgets
+from shared.types import Pose
+
+
+class ParticlesVisualization:
+    def __init__(
+        self,
+        states: list[np.ndarray],
+        drone_poses: Sequence[Pose] | Pose,
+        video: Sequence[np.ndarray] | None = None,
+    ):
+        """
+        Parameters
+        ----------
+        states : list[np.ndarray]
+            List of particle states at each timestep.
+        drone_poses : Sequence[Pose] | Pose | None
+            The poses of the drone at each timestep. If a single pose is provided,
+            it will be used for all iterations.
+        """
+        self._states = states
+
+        if isinstance(drone_poses, Pose):
+            self._drone_poses: Sequence[Pose] = list(repeat(drone_poses, len(states)))
+        elif isinstance(drone_poses, Sequence):
+            self._drone_poses: Sequence[Pose] = drone_poses
+
+        self._video = video
+
+        self._fig = plt.figure()
+
+        self._ax1 = self._fig.add_axes((0, 0, 1, 0.8), projection="3d")  # 3D plot axis
+
+        self._ax2 = self._fig.add_axes((0.1, 0.85, 0.8, 0.1))  # slider axis
+        self._slider = widgets.Slider(
+            self._ax2, "Time Step", 0, len(self._states) - 1, valinit=0
+        )
+
+        if self._video is not None:
+            self._ax3 = self._fig.add_axes((0.1, 0.7, 0.8, 0.15))  # camera frame axis
+            self._ax3.set_title("Camera Frame")
+        else:
+            self._ax3 = None
+
+        self._slider.on_changed(self._update_fig)
+
+    @staticmethod
+    def from_npy_dir(npy_dir: Path):
+        """
+        Creates a ParticlesVisualization from a directory containing .npy files.
+        This is useful to use in conjunction with `ArrayLogger`.
+
+        The files must be ordered by name.
+
+        Parameters
+        ----------
+        npy_dir : Path
+            Directory containing .npy files with particle states.
+        """
+        states = []
+        for file in sorted(npy_dir.glob("*.npy")):
+            states.append(np.load(file))
+
+        drone_poses = [Pose.identity()] * len(states)
+
+        return ParticlesVisualization(states=states, drone_poses=drone_poses)
+
+    def _update_fig(self, iteration: float):
+        if self._states is None:
+            raise RuntimeError("Simulation must be run before visualization")
+
+        iteration = round(iteration)
+        if iteration < 0:
+            iteration = 0
+
+        self._ax1.clear()
+        self._ax1.set_title(f"Iteration {iteration}")
+
+        particles = self._states[iteration]
+        self._ax1.scatter(particles[:, 0], particles[:, 1], particles[:, 2])
+
+        drone_pose = self._drone_poses[iteration]
+        direction = drone_pose.rotation.apply([0, 0, 1])
+
+        arrow_length = torch.mean(torch.tensor(particles).norm(dim=1)).item() / 10
+        self._ax1.quiver(
+            *drone_pose.position, *direction, length=arrow_length, color="r"
+        )
+
+        if self._video is not None and self._ax3 is not None:
+            self._ax3.clear()
+            self._ax3.imshow(self._video[iteration])
+            self._ax3.axis("off")
+
+    def visualize(self):
+        self._update_fig(0)
+        plt.show()
+
+
+if __name__ == "__main__":
+    visualizer = ParticlesVisualization.from_npy_dir(
+        Path("~/uavf_2025/logs/nvme/2025-06-12_13-37/obstacles/particles").expanduser()
+    )
+    visualizer.visualize()

uavf_2025/perception/lib/util.py

@@ -35,6 +35,8 @@ def format(self, record):
 def create_console_logger(
     logs_path: Path, name: str, capture_all_logs: bool = True
 ) -> logging.Logger:
+    logs_path.mkdir(parents=True, exist_ok=True)  # Ensure logs directory exists
+
     logger: logging.Logger = logging.getLogger(name)
     logger.setLevel(logging.DEBUG)  # Set the overall logging level