map.py

import yaml, cv2
from pathlib import Path
import numpy as np

type Number = float | int
type Point2D = tuple[Number, Number]


def project(x: Number, x0: Number, x1: Number) -> float:
    return (x - x0) / (x1 - x0)


def loadPGM(path: str | Path) -> np.ndarray:
    path = Path(path)
    with path.open("rb") as f:
        assert f.readline().strip() == b"P5"
        # Read width and height
        while True:
            line = f.readline()
            if line.startswith(b"#"):
                continue
            else:
                w, h = [int(x) for x in line.strip().split()]
                break
        v_max = int(f.readline().strip())
        assert v_max <= 255
        data = f.read()
        img = np.frombuffer(data, dtype=np.uint8).reshape((h, w))
        return img


class Map:
    @classmethod
    def load(cls, path: str | Path):
        img_path = Path(path)
        yml_path = img_path.with_suffix(".yaml")
        img = (
            loadPGM(img_path)
            if img_path.suffix == ".pgm"
            else cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE)
        )
        if not yml_path.exists():
            raise ValueError(f"YAML file not found for map: {yml_path}")
        yml: dict = yaml.safe_load(yml_path.read_text())
        # Validate img
        if img is None:
            raise ValueError(f"Failed to load image from {img_path}")
        img = img.astype(np.float32) / 255.0
        # Calculate tl and br from resolution and origin
        res = float(yml["resolution"])
        t0 = yml.get("free_thresh", 0.2)
        t1 = yml.get("occupied_thresh", 0.65)
        img[img < t0] = 0.0
        img[img > t1] = 1.0
        tl: Point2D = yml["origin"][:2]  # [x, y, theta]
        return tl, res, img

    def __init__(
        self,
        tl: Point2D,
        br: Point2D | None = None,
        canvas: np.ndarray | None = None,
        resolution: float = 0.01,
    ):
        """
        Initialize a Map from:
        1. tl + br + resolution
        2. tl + canvas + resolution
        """
        if canvas is None:
            if br is None:
                raise ValueError("Either br or canvas must be provided.")
            h = int(np.ceil((br[1] - tl[1]) / resolution))
            w = int(np.ceil((br[0] - tl[0]) / resolution))
            canvas = np.ones((h, w), dtype=np.float32)
        elif br is not None:
            raise ValueError("Cannot provide both br and canvas.")
        self.tl = tl
        self.resolution = resolution
        self.canvas = canvas

    @property
    def br(self) -> Point2D:
        h, w = self.canvas.shape
        x1 = self.tl[0] + w * self.resolution
        y1 = self.tl[1] + h * self.resolution
        return (x1, y1)

    def meter2px(self, p: Point2D) -> Point2D:
        h, w = self.canvas.shape
        a, b = p
        x0, y0 = self.tl
        x1, y1 = self.br
        return w * (project(b, x0, x1)), h * project(-a, y0, y1)

    def line(
        self, p0: Point2D, p1: Point2D, radius: float, img: np.ndarray | None = None
    ):
        """
        Draw a rounded line on the map canvas or the provided img.
        """
        if img is None:
            img = self.canvas
        # First draw the line
        x0, y0 = int(round(p0[0])), int(round(p0[1]))
        x1, y1 = int(round(p1[0])), int(round(p1[1]))
        r = int(round(radius / self.resolution))
        color = 255 if img.dtype == np.uint8 else 1.0
        cv2.line(img, (x0, y0), (x1, y1), color=color, thickness=2 * r)
        # Then draw circles at the endpoints
        cv2.circle(img, (x0, y0), r, color=color, thickness=-1)
        cv2.circle(img, (x1, y1), r, color=color, thickness=-1)

    def render(self, c0: int = 192, c1: int = 128, c2: int = 255, t: float = 0.1):
        t = int(round(t / self.resolution))
        img = np.zeros((*self.canvas.shape, 3), dtype=np.uint8)
        img[self.canvas == 0.0] = [c0, c0, c0]

        # Create an outline outside white areas (value 1.0)
        white_mask = (self.canvas == 1.0).astype(np.uint8)
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * t + 1, 2 * t + 1))
        dilated = cv2.dilate(white_mask, kernel, iterations=1)
        outline = dilated - white_mask

        # Apply colors: outline gets c1, white areas get c2
        img[outline == 1] = [c1, c1, c1]
        img[self.canvas == 1.0] = [c2, c2, c2]
        return img


if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser(description="Test Map loading and drawing")
    parser.add_argument(
        "map_path", type=Path, help="Path to the map image (.pgm or other)"
    )
    args = parser.parse_args()
    tl, res, img = Map.load(args.map_path)
    print(f"Loaded map from {args.map_path}")
    print(f"Top-left: {tl}, Resolution: {res}, Shape: {img.shape}")
    m = Map(tl=tl, canvas=img, resolution=res)
    # Render map
    cv2.imshow("Map", m.render()[::-1])
    cv2.imshow("Map (raw)", m.canvas[::-1])  # --- IGNORE ---
    cv2.waitKey(0)
    cv2.destroyAllWindows()