119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
# SLAM_MAP.py
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import numpy as np
|
|
|
|
|
|
def load_slam_config() -> dict:
|
|
import json, os
|
|
from pathlib import Path
|
|
|
|
cfg_path = os.environ.get("SLAM_CONFIG", "").strip()
|
|
p = Path(cfg_path) if cfg_path else (Path(__file__).resolve().parent / "SLAM_Config.json")
|
|
|
|
if not p.exists():
|
|
raise FileNotFoundError(f"Missing config file: {p}")
|
|
|
|
text = p.read_text(encoding="utf-8").strip()
|
|
if not text:
|
|
raise RuntimeError("SLAM_Config.json is empty.")
|
|
|
|
return json.loads(text)
|
|
|
|
|
|
@dataclass
|
|
class MapConfig:
|
|
display_voxel: float
|
|
save_voxel: float
|
|
data_folder: str
|
|
save_extension: str
|
|
|
|
@staticmethod
|
|
def from_config_file() -> "MapConfig":
|
|
cfg = load_slam_config()
|
|
maps_dir = str(cfg["app"]["maps_dir"])
|
|
return MapConfig(
|
|
display_voxel=float(cfg["map"]["display_voxel"]),
|
|
save_voxel=float(cfg["map"]["save_voxel"]),
|
|
data_folder=maps_dir,
|
|
save_extension=str(cfg["map"]["save_extension"]),
|
|
)
|
|
|
|
|
|
def _voxel_downsample_numpy(points: np.ndarray, voxel: float) -> np.ndarray:
|
|
if points is None or len(points) == 0:
|
|
return np.zeros((0, 3), dtype=np.float32)
|
|
if voxel <= 0:
|
|
return points.astype(np.float32, copy=False)
|
|
|
|
keys = np.floor(points / voxel).astype(np.int32)
|
|
uniq, inv = np.unique(keys, axis=0, return_inverse=True)
|
|
|
|
out = np.zeros((len(uniq), 3), dtype=np.float64)
|
|
cnt = np.zeros((len(uniq),), dtype=np.int64)
|
|
np.add.at(out, inv, points.astype(np.float64, copy=False))
|
|
np.add.at(cnt, inv, 1)
|
|
out /= np.maximum(cnt[:, None], 1)
|
|
return out.astype(np.float32)
|
|
|
|
|
|
class StableMapLayer:
|
|
def __init__(self, cfg: MapConfig | None = None):
|
|
self.cfg = cfg if cfg is not None else MapConfig.from_config_file()
|
|
Path(self.cfg.data_folder).mkdir(parents=True, exist_ok=True)
|
|
self._points = np.zeros((0, 3), dtype=np.float32)
|
|
|
|
def reset(self):
|
|
self._points = np.zeros((0, 3), dtype=np.float32)
|
|
|
|
def set_points(self, pts: np.ndarray):
|
|
self._points = pts.astype(np.float32, copy=True) if pts is not None else np.zeros((0, 3), dtype=np.float32)
|
|
|
|
def get_points(self) -> np.ndarray:
|
|
return self._points
|
|
|
|
def get_display_points(self) -> np.ndarray:
|
|
return _voxel_downsample_numpy(self._points, float(self.cfg.display_voxel))
|
|
|
|
def get_save_points(self) -> np.ndarray:
|
|
return _voxel_downsample_numpy(self._points, float(self.cfg.save_voxel))
|
|
|
|
def export_map(self, base_name: str) -> str:
|
|
"""
|
|
Save stable map into maps_dir (no subfolders).
|
|
"""
|
|
import open3d as o3d
|
|
|
|
pts = self.get_save_points()
|
|
if pts is None or len(pts) < 20:
|
|
raise RuntimeError("Not enough stable points to save.")
|
|
|
|
base = (base_name or "").strip() or "map_robot"
|
|
ext = self.cfg.save_extension.strip() or ".ply"
|
|
if base.endswith(ext):
|
|
base = base[: -len(ext)]
|
|
|
|
filename = f"{base}{ext}"
|
|
n = 1
|
|
folder = Path(self.cfg.data_folder)
|
|
while (folder / filename).exists():
|
|
filename = f"{base}({n}){ext}"
|
|
n += 1
|
|
|
|
full = str(folder / filename)
|
|
pcd = o3d.geometry.PointCloud()
|
|
pcd.points = o3d.utility.Vector3dVector(pts.astype(np.float64, copy=False))
|
|
o3d.io.write_point_cloud(full, pcd)
|
|
return full
|
|
|
|
def load_map(self, filepath: str) -> np.ndarray:
|
|
import open3d as o3d
|
|
pcd = o3d.io.read_point_cloud(filepath)
|
|
pts = np.asarray(pcd.points).astype(np.float32)
|
|
self.set_points(pts)
|
|
return pts
|