Spaces:
Sleeping
Sleeping
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| 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) # add background | |
| 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]) | |