occupancy_grid.py

#! /usr/bin/env python3
# Copyright 2025 Arjo Chakravarty
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
 
"""
This is a Python3 API for costmap 2d messages from the stack.
 
It provides the basic conversion, get/set,
and handling semantics found in the costmap 2d C++ API.
"""
 
 
from nav_msgs.msg import OccupancyGrid
import numpy as np
 
 
class PyOccupancyGrid:
    """
    PyOccupancyGrid.
 
    Occupancy grid Python3 API for OccupancyGrids to populate from published messages
    """
 
    def __init__(self, occupancy_map: OccupancyGrid):
        """
        Initialize costmap2D.
 
        Args
        ----
            occupancy_map (OccupancyGrid): 2D OccupancyGrid Map
 
        Returns
        -------
            None
 
        """
        self.size_xsize_x = occupancy_map.info.width
        self.size_ysize_y = occupancy_map.info.height
        self.resolutionresolution = occupancy_map.info.resolution
        self.origin_xorigin_x = occupancy_map.info.origin.position.x
        self.origin_yorigin_y = occupancy_map.info.origin.position.y
        self.global_frame_idglobal_frame_id = occupancy_map.header.frame_id
        self.costmap_timestampcostmap_timestamp = occupancy_map.header.stamp
        # Extract costmap
        self.costmapcostmap = np.array(occupancy_map.data, dtype=np.int8)
 
    def getSizeInCellsX(self):
        """Get map width in cells."""
        return self.size_xsize_x
 
    def getSizeInCellsY(self):
        """Get map height in cells."""
        return self.size_ysize_y
 
    def getSizeInMetersX(self):
        """Get x axis map size in meters."""
        return self.size_xsize_x * self.resolutionresolution
 
    def getSizeInMetersY(self):
        """Get y axis map size in meters."""
        return self.size_ysize_y * self.resolutionresolution
 
    def getOriginX(self):
        """Get the origin x axis of the map [m]."""
        return self.origin_xorigin_x
 
    def getOriginY(self):
        """Get the origin y axis of the map [m]."""
        return self.origin_yorigin_y
 
    def getResolution(self):
        """Get map resolution [m/cell]."""
        return self.resolutionresolution
 
    def getGlobalFrameID(self):
        """Get global frame_id."""
        return self.global_frame_idglobal_frame_id
 
    def getCostmapTimestamp(self):
        """Get costmap timestamp."""
        return self.costmap_timestampcostmap_timestamp
 
    def getCostXY(self, mx: int, my: int):
        """
        Get the cost of a cell in the costmap using map coordinate XY.
 
        Args
        ----
            mx (int): map coordinate X to get cost
            my (int): map coordinate Y to get cost
 
        Returns
        -------
            np.int8: cost of a cell
 
        """
        return np.int8(self.costmapcostmap[self.getIndexgetIndex(mx, my)])
 
    def getCostIdx(self, index: int):
        """
        Get the cost of a cell in the costmap using Index.
 
        Args
        ----
            index (int): index of cell to get cost
 
        Returns
        -------
            np.int8: cost of a cell
 
        """
        return np.int8(self.costmapcostmap[index])
 
    def setCost(self, mx: int, my: int, cost: np.int8):
        """
        Set the cost of a cell in the costmap using map coordinate XY.
 
        Args
        ----
            mx (int): map coordinate X to get cost
            my (int): map coordinate Y to get cost
            cost (np.int8): The cost to set the cell
 
        Returns
        -------
            None
 
        """
        self.costmapcostmap[self.getIndexgetIndex(mx, my)] = cost
 
    def mapToWorld(self, mx: int, my: int):
        """
        Get the world coordinate XY using map coordinate XY.
 
        Args
        ----
            mx (int): map coordinate X to get world coordinate
            my (int): map coordinate Y to get world coordinate
 
        Returns
        -------
            tuple of float: wx, wy
            wx (float) [m]: world coordinate X
            wy (float) [m]: world coordinate Y
 
        """
        wx = self.origin_xorigin_x + (mx + 0.5) * self.resolutionresolution
        wy = self.origin_yorigin_y + (my + 0.5) * self.resolutionresolution
        return (wx, wy)
 
    def worldToMapValidated(self, wx: float, wy: float):
        """
        Get the map coordinate XY using world coordinate XY.
 
        Args
        ----
            wx (float) [m]: world coordinate X to get map coordinate
            wy (float) [m]: world coordinate Y to get map coordinate
 
        Returns
        -------
            (None, None): if coordinates are invalid
            tuple of int: mx, my (if coordinates are valid)
            mx (int): map coordinate X
            my (int): map coordinate Y
 
        """
        if wx < self.origin_xorigin_x or wy < self.origin_yorigin_y:
            return (None, None)
        mx = int((wx - self.origin_xorigin_x) // self.resolutionresolution)
        my = int((wy - self.origin_yorigin_y) // self.resolutionresolution)
        if mx < self.size_xsize_x and my < self.size_ysize_y:
            return (mx, my)
        return (None, None)
 
    def getIndex(self, mx: int, my: int):
        """
        Get the index of the cell using map coordinate XY.
 
        Args
        ----
            mx (int): map coordinate X to get Index
            my (int): map coordinate Y to get Index
 
        Returns
        -------
            int: The index of the cell
 
        """
        return my * self.size_xsize_x + mx