visualization.py

import numpy as np
import cv2


_COLORMAPS = {
    "viridis": cv2.COLORMAP_VIRIDIS,
    "jet": cv2.COLORMAP_JET,
    "hot": cv2.COLORMAP_HOT,
    "inferno": cv2.COLORMAP_INFERNO,
    "plasma": cv2.COLORMAP_PLASMA,
}


class MapVisualizer:
    def __init__(self, config: dict):
        self.layers = config.get("layers", ["elevation"])
        cmap_name = config.get("colormap", "viridis")
        self.colormap = _COLORMAPS.get(cmap_name, cv2.COLORMAP_VIRIDIS)
        self.window_size = config.get("window_size", 512)

    def update(self, elevation_map, position):
        cell_n = elevation_map.param.cell_n
        data = np.zeros((cell_n - 2, cell_n - 2), dtype=np.float32)
        panels = []

        for layer in self.layers:
            try:
                elevation_map.get_map_with_name_ref(layer, data)
            except Exception:
                continue

            layer_data = data.copy()
            nan_mask = np.isnan(layer_data)

            # Normalize to [0, 255]
            valid = layer_data[~nan_mask]
            if valid.size > 0:
                vmin, vmax = valid.min(), valid.max()
                if vmax > vmin:
                    norm = (layer_data - vmin) / (vmax - vmin)
                else:
                    norm = np.zeros_like(layer_data)
            else:
                norm = np.zeros_like(layer_data)

            norm = np.clip(norm, 0.0, 1.0)
            gray = (norm * 255).astype(np.uint8)
            colored = cv2.applyColorMap(gray, self.colormap)
            colored[nan_mask] = 0  # black for unknown cells

            # Add layer label
            cv2.putText(colored, layer, (5, 20), cv2.FONT_HERSHEY_SIMPLEX,
                        0.6, (255, 255, 255), 1, cv2.LINE_AA)

            # Resize to window_size
            colored = cv2.resize(colored, (self.window_size, self.window_size),
                                 interpolation=cv2.INTER_NEAREST)
            panels.append(colored)

        if not panels:
            return

        display = np.hstack(panels)
        cv2.imshow("Elevation Map", display)
        cv2.waitKey(1)

    def close(self):
        cv2.destroyAllWindows()