coverage.py

from typing import Iterable
import numpy as np
import cv2
from pathlib import Path

from map import Map

from config import (
    RADIUS as R,
    MAP_RESOLUTION as MR,
    TIME_RESOLUTION as TR,
    DISPLAY,
    LOAD_MAP,
    SAVE_MAP,
)

type Point2D = tuple[float, float]


def timer(func):
    import time

    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"Function {func.__name__} took {end - start:.6f} seconds")
        return result

    return wrapper


class Trajectory(Map):
    def __init__(self, input: Iterable[str], load_map: str | Path | None = None):
        data: list[list[float]] = []
        for line in input:
            line = line.strip()
            if len(line) == 0 or line.startswith("#"):
                continue
            try:
                data.append([float(x) for x in line.split(",")])
            except ValueError:
                continue
        data.sort(key=lambda x: x[0])
        if len(data) == 0:
            raise ValueError("No valid trajectory points found.")
        # ts, x, y, r
        self.data = np.array(data, dtype=np.float64)
        self.data[:, 0] -= self.data[0, 0]  # Normalize time to start from 0
        self.time_range = (self.data[0, 0], self.data[-1, 0])
        self.idx = 0
        self.t = self.time_range[0]
        # Construct map
        if load_map is not None:
            tl, res, img = Map.load(load_map)
            super().__init__(tl=tl, canvas=img, resolution=res)
        else:
            x0 = np.min(self.data[:, 1])
            x1 = np.max(self.data[:, 1])
            y0 = np.min(self.data[:, 2])
            y1 = np.max(self.data[:, 2])
            tl = (x0 - R, y0 - R)
            br = (x1 + R, y1 + R)
            super().__init__(tl=tl, br=br, resolution=MR)
        h, w = self.canvas.shape[:2]
        self.trj_canvas = np.zeros((h, w), dtype=np.float32)
        print("Map", self.tl, self.br)

    def step(self, dt: float) -> bool:
        self.t += dt
        flag_updated = False
        while self.idx < len(self.data) - 1 and self.data[self.idx + 1][0] <= self.t:
            x0, y0 = self.data[self.idx][1:3]
            self.idx += 1
            x1, y1 = self.data[self.idx][1:3]
            # Paint the canvas along the line from (x0, y0) to (x1, y1) with radius R
            p0 = self.meter2px((x0, y0))
            p1 = self.meter2px((x1, y1))
            self.line(p0, p1, R, self.trj_canvas)
            self._blurred_cache = None  # Invalidate cache
            flag_updated = True
        return flag_updated

    _blurred_cache: np.ndarray | None = None

    @property
    def blurred(self) -> np.ndarray:
        if self._blurred_cache is None:
            self._blurred_cache = cv2.GaussianBlur(
                self.trj_canvas, (0, 0), sigmaX=R / self.resolution
            )
            # self._blurred_cache /= sigma_px * np.sqrt(2 * np.pi)
            self._blurred_cache = np.clip(self._blurred_cache, 0.0, 1.0)
        return self._blurred_cache

    @property
    def coverage(self) -> float:
        area = self.resolution**2
        # gaussian blur to smooth the trajectory coverage
        covered = float(np.sum(self.blurred * self.canvas))
        return covered * area

    def renderTrj(self) -> np.ndarray:
        map = self.render()
        free = self.canvas == 1.0  # Free space mask
        # Heatmap the trajectory
        trj_color = np.stack(
            [
                self.blurred,  # B
                np.zeros_like(self.blurred),  # G
                1 - self.blurred,  # R
            ],
            axis=-1,
        )
        # Overlay
        trj_color = (trj_color * 255).astype(np.uint8)
        map[free] = map[free] / 2 + trj_color[free] / 2
        return map[::-1]  # Flip vertically for display


if __name__ == "__main__":
    try:
        import sys

        trj = Trajectory(sys.stdin, LOAD_MAP)
        while trj.step(TR):
            print(f"{trj.t},{trj.coverage:.6f}")
            if DISPLAY:
                cv2.imshow("Coverage", trj.renderTrj())
                key = cv2.waitKey(1)
                if key == 27 or key == ord("q"):  # ESC key or 'q'
                    break
        if DISPLAY:
            while cv2.waitKey(1) < 0:
                pass
            cv2.destroyAllWindows()
        if SAVE_MAP:
            cv2.imwrite(str(SAVE_MAP), trj.renderTrj())
    except KeyboardInterrupt:
        pass