Add Lidar2D class for future extension

README.md

@@ -6,10 +6,10 @@
 [![Build](https://github.com/prina404/raycast2D/actions/workflows/build.yml/badge.svg)](https://github.com/prina404/raycast2D/actions/workflows/build.yml)
 [![Tests](https://github.com/prina404/raycast2D/actions/workflows/tests.yml/badge.svg)](https://github.com/prina404/raycast2D/actions/workflows/tests.yml)
 
-`raycast2D` is a fast, single-core 2D raycasting implementation written in C with a small Python API.
+`raycast2D` is a fast, single-thread 2D raycasting implementation written in C with a small Python API.
 It operates on NumPy occupancy grids / binary images (free cells are non-zero; occupied cells are `0`) and computes ray intersections using Bresenham line algorithm.
 
-The core package depends only on `numpy`. Optional extras are provided for the interactive demo and development.
+The core package depends only on `numpy`.
 
 ![](media/raycast_compressed.gif)
 
@@ -18,49 +18,83 @@ The core package depends only on `numpy`. Optional extras are provided for the i
 ```bash
 $ pip install raycast2D
 ```
-### Basic raycast on an image
+### Basic `cast` function usage
 
 ```python
 import numpy as np
-from raycast2D import cast
 from PIL import Image
-import matplotlib.pyplot as plt
+from raycast2D import cast
 
 img = Image.open("<path/to/img.png>")
-img_array = np.array(img)
-img_array = img_array[:, :, 0].astype(np.uint8)  # Use single channel
-# img_array = img_array[img_array < 100] = 0  # threshold obstacles if needed
+img_array = np.array(img)[:, :, 0] # Use single channel
+pose = (250, 250, 0.5)  # x, y, yaw in radians 
+
+rays = cast(img_array, pose, num_rays=360, ray_length=500)
+
+print(rays[:5])
+# Expected output (example):
+# [[x0 y0]
+#  [x1 y1]
+#  [x2 y2]
+#  [x3 y3]
+#  [x4 y4]]
+```
+
+### Using the `Lidar2D` class
 
-rays = cast(img_array, pose=(250, 250), num_rays=360, ray_length=500)
-# rays is an array of shape (N, 2) where each row contains the (x, y) coordinates of the ray collisions
+```python
+from raycast2D import Lidar2D
+# [...] 
+lidar = Lidar2D(num_rays=360, FOV=360, ray_length=500)
+lidar.set_map(img_array)
 
-plt.imshow(img_array, cmap='gray')
-plt.scatter([250], [250], c='green', s=10)
-plt.scatter(rays[:, 0], rays[:, 1], c='blue', s=1)
+rays = lidar.scan((250, 250, 0.5))
+
+print(rays[:5])
+# same output as before...
+```
+
+### Plotting example
+
+```python
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+from raycast2D import cast
+
+img = Image.open("media/lab_intel.png")
+img_array = np.array(img)[:, :, 0] # Use single channel
+
+pose = (350, 350, -1.0)
+rays = cast(img_array, pose=pose, num_rays=200, FOV=90, ray_length=250)
+
+plt.imshow(img_array, cmap="gray")
+plt.scatter([pose[0]], [pose[1]], c="green", s=10)
+plt.scatter(rays[:, 0], rays[:, 1], c="blue", s=1)
 for ray in rays:
-    plt.plot([250, ray[0]], [250, ray[1]], c='red', linewidth=0.5, alpha=0.3)
+    plt.plot([pose[0], ray[0]], [pose[1], ray[1]], c="red", linewidth=0.5, alpha=0.3)
 plt.show()
 ```
 
+![](media/example.png)
+
 ## Performance
 
 The benchmark script in `test/benchmark.py` runs a small set of examples and prints throughput in rays/s.
 Example results (tested on an i7-9700k):
 
 |  Map size | Ray length | Rays per call | Mean time (ms) | Throughput (rays/s) |
 | --------: | ---------: | ------------: | -------------: | ------------------: |
-|   512×512 |       2000px |          5000 |         0.2904 |          17,216,880 |
-| 4096×4096 |       2000px |          5000 |         1.7626 |           2,836,661 |
-| 8192×8192 |       2000px |          5000 |         7.3452 |             680,719 |
+|   512×512 |     2000px |          5000 |         0.2490 |          20,083,924 |
+| 4096×4096 |     2000px |          5000 |         0.2848 |          17,556,549 |
+| 8192×8192 |     2000px |          5000 |         0.2571 |          19,446,324 |
 
 To reproduce on your machine:
 
 ```bash
 $ python3 test/benchmark.py
 ```
 
-
-
 ## Interactive demo
 
 The repository includes an interactive `pygame` demo that raycasts from the current mouse position.

src/raycast2D/__init__.py

@@ -1,3 +1,3 @@
-from .raycast2D import cast
+from .raycast2D import cast, Lidar2D
 
-__all__ = ["cast"]
+__all__ = ["cast", "Lidar2D"]

src/raycast2D/raycast2D.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from . import raycaster
 import numpy as np
 from numpy.typing import NDArray
@@ -6,6 +7,94 @@
 import warnings
 
 
+class Lidar2D:
+
+    def __init__(
+        self,
+        num_rays: int = 1000,
+        FOV: int = 360,
+        ray_length: int = 500,
+        occupancy_grid: NDArray | None = None,
+        obstacle_threshold: int = 0,
+    ) -> None:
+        """
+        Initialize the Lidar2D raycaster.
+        Args:
+            num_rays: Number of rays to cast (samples uniformly over the FOV).
+            FOV: Field of view in degrees.
+            ray_length: Max ray length in pixels.
+            occupancy_grid: Optional occupancy grid to set at initialization.
+            obstacle_threshold: Cells with values less than or equal to this are treated as occupied.
+        """
+        self.angles = np.linspace(
+            -math.radians(FOV) / 2.0,
+            math.radians(FOV) / 2.0,
+            num_rays,
+        )
+        self.ray_length = ray_length
+        self.obstacle_threshold = obstacle_threshold
+        self._map = None
+        if occupancy_grid is not None:
+            self.set_map(occupancy_grid)
+
+    def scan(
+        self,
+        pose: tuple[int, int] | tuple[int, int, float],
+        image: Optional[NDArray] = None,
+        only_true_collisions: bool = True
+    ) -> NDArray[np.uint32]:
+        """
+        Perform a LiDAR scan from the given pose on the occupancy image.
+
+        Args:
+            pose: ``(x, y)`` or ``(x, y, yaw)`` where ``yaw`` is in radians.
+            image: Optional occupancy grid of shape ``(H, W)`` or ``(H, W, 1)`` with an integer dtype.
+                The raycaster treats values less than or equal to ``obstacle_threshold`` as occupied;
+                the start cell at ``pose`` must be free.
+                NOTE: the provided image overrides any occupancy grid set at initialization.
+            only_true_collisions: If True, return only rays that actually hit an obstacle
+                (LiDAR-style). If False, return endpoints for all rays.
+
+        Returns:
+            An array of shape ``(N, 2)`` with dtype ``uint32`` containing ``(x, y)`` collision coordinates.
+        """
+        if self._map is None and image is None:
+            raise ValueError(
+                "No occupancy grid provided at initialization or scan time.")
+
+        if image is not None :
+            _map = (_check_map(image) > self.obstacle_threshold).astype(np.uint8)
+        else:
+            _map = self._map
+
+        _check_pose(_map, pose)
+
+        p = _POSE(*pose)  # type: ignore
+        x_rays = np.round(p.x + self.ray_length *
+                          np.cos(self.angles + p.yaw)).astype(np.uint32)
+        y_rays = np.round(p.y + self.ray_length *
+                          np.sin(self.angles + p.yaw)).astype(np.uint32)
+
+        rays = np.column_stack((x_rays, y_rays, np.zeros_like(y_rays)))
+
+        # raycaster modifies the rays array inplace to include collisions
+        raycaster.raycast(_map, rays, p.x, p.y)
+
+        if only_true_collisions:    # LiDaR-style output
+            return rays[rays[:, -1] == 1][:, :-1]
+        return rays[:, :-1]
+
+    def set_map(self, occupancy_grid: NDArray) -> None:
+        """
+        Set or update the occupancy grid used by the raycaster.
+
+        Args:
+            occupancy_grid: Occupancy grid of shape ``(H, W)`` or ``(H, W, 1)`` with an integer dtype.
+        """
+        self._map = _check_map(occupancy_grid)
+        self._map = (self._map > self.obstacle_threshold).astype(np.uint8)
+
+
 def cast(
     image: NDArray,
     pose: tuple[int, int] | tuple[int, int, float],
@@ -31,62 +120,72 @@ def cast(
         An array of shape ``(N, 2)`` with dtype ``uint32`` containing ``(x, y)`` collision coordinates.
         If ``only_true_collisions=True``, ``N`` is the number of hits; otherwise ``N == num_rays``.
     """
-    _check_args(image, pose)
-    if len(image.shape) == 3:   # take only the first channel
-        image = image[:, :, 0]
-    if image.dtype != np.uint8:
-        warnings.warn(f"Converting image of dtype {image.dtype} to uint8 for raycasting.")
-        image = image.astype(np.uint8)
-
-    p = _POSE(*pose) # type: ignore
+    image = _check_map(image)
+    _check_pose(image, pose)
+
+    p = _POSE(*pose)  # type: ignore
     half_fov = math.radians(FOV) / 2.0
     start_angle = p.yaw - half_fov
     end_angle = p.yaw + half_fov
 
     theta = np.linspace(start_angle, end_angle, num_rays)
     x_rays = np.round(p.x + ray_length * np.cos(theta)).astype(np.uint32)
     y_rays = np.round(p.y + ray_length * np.sin(theta)).astype(np.uint32)
-    
+
     rays = np.column_stack((x_rays, y_rays, np.zeros_like(y_rays)))
-    
+
     # raycaster modifies the rays array inplace to include collisions
     raycaster.raycast(image, rays, p.x, p.y)
-    
+
     if only_true_collisions:    # LiDaR-style output
-        return rays[rays[:, 2] == 1][:, :2]
-    return rays[:, :2]
+        return rays[rays[:, -1] == 1][:, :-1]
+    return rays[:, :-1]
 
 
 @dataclass
 class _POSE:
     x: int
     y: int
-    yaw: float = -math.pi / 2
-
-
-def _check_args(img: NDArray, pose: tuple) -> None:
-    if not np.issubdtype(img.dtype, np.integer):
-        raise ValueError(
-            f"Received an array with dtype {img.dtype}, The only supported dtypes are subtypes of np.integral.")
+    yaw: float = 0.0
+
+    def __post_init__(self):
+        self.yaw -= math.pi / 2.0  # adjust so yaw=0 points upwards
 
-    shape = img.shape
-    if len(shape) < 2 or (len(shape) > 2 and shape[2] > 1):
-        raise ValueError(
-            f"Received an array of shape {shape}. The only supported shapes are (H x W), (H x W x 1)."
-        )
+def _check_pose(img: NDArray[np.uint8], pose: tuple) -> None:
     if len(pose) < 2 or len(pose) > 3:
         raise ValueError(
             f"Received a pose tuple {pose}. The only supported pose formats are (x, y) and (x, y, yaw).")
 
     x, y = pose[0], pose[1]
-    H, W = shape[0], shape[1]
+    H, W = img.shape[0], img.shape[1]
     if x < 0 or x >= W:
         raise IndexError(
             f"x coordinate ({x}) if out of bounds for array with width {W}.")
     if y < 0 or y >= H:
         raise IndexError(
             f"y coordinate ({y}) if out of bounds for array with height {H}.")
-        
+
     if img[y, x] == 0:
         raise ValueError(
             f"The pose {pose} corresponds to an occupied cell, cannot raycast from an occupied cell.")
+
+
+def _check_map(img: NDArray) -> NDArray[np.uint8]:
+    if not np.issubdtype(img.dtype, np.integer):
+        raise ValueError(
+            f"Received an array with dtype {img.dtype}, The only supported dtypes are subtypes of np.integral.")
+
+    shape = img.shape
+    if len(shape) < 2 or (len(shape) > 2 and shape[2] > 1):
+        raise ValueError(
+            f"Received an array of shape {shape}. The only supported shapes are (H x W), (H x W x 1)."
+        )
+
+    if len(img.shape) == 3:   # take only the first channel
+        img = img[:, :, 0]
+    if img.dtype != np.uint8:
+        warnings.warn(
+            f"Converting image of dtype {img.dtype} to uint8 for raycasting.")
+        img = img.astype(np.uint8)
+
+    return img