|
|
|
|
|
from typing import Dict, List |
|
|
|
import cv2 |
|
import numpy as np |
|
import torch |
|
|
|
from utils.geo import BoundaryBox |
|
from .data import MapArea, MapLine, MapNode |
|
from .parser import Groups |
|
|
|
|
|
class Canvas: |
|
def __init__(self, bbox: BoundaryBox, ppm: float): |
|
self.bbox = bbox |
|
self.ppm = ppm |
|
self.scaling = bbox.size * ppm |
|
self.w, self.h = np.ceil(self.scaling).astype(int) |
|
self.clear() |
|
|
|
def clear(self): |
|
self.raster = np.zeros((self.h, self.w), np.uint8) |
|
|
|
def to_uv(self, xy: np.ndarray): |
|
xy = self.bbox.normalize(xy) |
|
xy[..., 1] = 1 - xy[..., 1] |
|
s = self.scaling |
|
if isinstance(xy, torch.Tensor): |
|
s = torch.from_numpy(s).to(xy) |
|
return xy * s - 0.5 |
|
|
|
def to_xy(self, uv: np.ndarray): |
|
s = self.scaling |
|
if isinstance(uv, torch.Tensor): |
|
s = torch.from_numpy(s).to(uv) |
|
xy = (uv + 0.5) / s |
|
xy[..., 1] = 1 - xy[..., 1] |
|
return self.bbox.unnormalize(xy) |
|
|
|
def draw_polygon(self, xy: np.ndarray): |
|
uv = self.to_uv(xy) |
|
cv2.fillPoly(self.raster, uv[None].astype(np.int32), 255) |
|
|
|
def draw_multipolygon(self, xys: List[np.ndarray]): |
|
uvs = [self.to_uv(xy).round().astype(np.int32) for xy in xys] |
|
cv2.fillPoly(self.raster, uvs, 255) |
|
|
|
def draw_line(self, xy: np.ndarray, width: float = 1): |
|
uv = self.to_uv(xy) |
|
cv2.polylines( |
|
self.raster, uv[None].round().astype(np.int32), False, 255, thickness=width |
|
) |
|
|
|
def draw_cell(self, xy: np.ndarray): |
|
if not self.bbox.contains(xy): |
|
return |
|
uv = self.to_uv(xy) |
|
self.raster[tuple(uv.round().astype(int).T[::-1])] = 255 |
|
|
|
|
|
def render_raster_masks( |
|
nodes: List[MapNode], |
|
lines: List[MapLine], |
|
areas: List[MapArea], |
|
canvas: Canvas, |
|
) -> Dict[str, np.ndarray]: |
|
all_groups = Groups.areas + Groups.ways + Groups.nodes |
|
masks = {k: np.zeros((canvas.h, canvas.w), np.uint8) for k in all_groups} |
|
|
|
for area in areas: |
|
canvas.raster = masks[area.group] |
|
outlines = area.outers + area.inners |
|
canvas.draw_multipolygon(outlines) |
|
if area.group == "building": |
|
canvas.raster = masks["building_outline"] |
|
for line in outlines: |
|
canvas.draw_line(line) |
|
|
|
for line in lines: |
|
canvas.raster = masks[line.group] |
|
canvas.draw_line(line.xy) |
|
|
|
for node in nodes: |
|
canvas.raster = masks[node.group] |
|
canvas.draw_cell(node.xy) |
|
|
|
return masks |
|
|
|
|
|
def mask_to_idx(group2mask: Dict[str, np.ndarray], groups: List[str]) -> np.ndarray: |
|
masks = np.stack([group2mask[k] for k in groups]) > 0 |
|
void = ~np.any(masks, 0) |
|
idx = np.argmax(masks, 0) |
|
idx = np.where(void, np.zeros_like(idx), idx + 1) |
|
return idx |
|
|
|
|
|
def render_raster_map(masks: Dict[str, np.ndarray]) -> np.ndarray: |
|
areas = mask_to_idx(masks, Groups.areas) |
|
ways = mask_to_idx(masks, Groups.ways) |
|
nodes = mask_to_idx(masks, Groups.nodes) |
|
return np.stack([areas, ways, nodes]) |
|
|
