import numpy as np from mmdet.datasets.builder import PIPELINES from shapely.geometry import LineString @PIPELINES.register_module(force=True) class PolygonizeLocalMapBbox(object): """Pre-Processing used by vectormapnet model. Args: canvas_size (tuple or list): bev feature size coord_dim (int): dimension of point's coordinate num_class (int): number of classes threshold (float): threshold for minimum bounding box size """ def __init__(self, canvas_size=(200, 100), coord_dim=2, num_class=3, threshold=6 / 200, ): self.canvas_size = np.array(canvas_size) self.num_class = num_class # for keypoints self.threshold = threshold self.coord_dim = coord_dim self.map_stop_idx = 0 self.coord_dim_start_idx = 1 def format_polyline_map(self, vectors): polylines, polyline_masks, polyline_weights = [], [], [] # quantilize each label's lines individually. for label, _lines in vectors.items(): for polyline in _lines: # and pad polyline. if label == 2: polyline_weight = evaluate_line(polyline).reshape(-1) else: polyline_weight = np.ones_like(polyline).reshape(-1) polyline_weight = np.pad( polyline_weight, ((0, 1),), constant_values=1.) polyline_weight = polyline_weight / polyline_weight.sum() # flatten and quantilized fpolyline = quantize_verts( polyline, self.canvas_size, self.coord_dim) fpolyline = fpolyline.reshape(-1) # reindex starting from 1, and add a zero stopping token(EOS), fpolyline = \ np.pad(fpolyline + self.coord_dim_start_idx, ((0, 1),), constant_values=0) fpolyline_msk = np.ones(fpolyline.shape, dtype=np.bool) polyline_masks.append(fpolyline_msk) polyline_weights.append(polyline_weight) polylines.append(fpolyline) polyline_map = polylines polyline_map_mask = polyline_masks polyline_map_weights = polyline_weights return polyline_map, polyline_map_mask, polyline_map_weights def format_keypoint(self, vectors): kps, kp_labels = [], [] qkps, qkp_masks = [], [] # quantilize each label's lines individually. for label, _lines in vectors.items(): for polyline in _lines: kp = get_bbox(polyline, self.threshold) kps.append(kp) kp_labels.append(label) gkp = kp # flatten and quantilized fkp = quantize_verts(gkp, self.canvas_size, self.coord_dim) fkp = fkp.reshape(-1) fkps_msk = np.ones(fkp.shape, dtype=np.bool) qkp_masks.append(fkps_msk) qkps.append(fkp) qkps = np.stack(qkps) qkp_msks = np.stack(qkp_masks) # format det kps = np.stack(kps, axis=0).astype(np.float32) * self.canvas_size kp_labels = np.array(kp_labels) # restrict the boundary kps[..., 0] = np.clip(kps[..., 0], 0.1, self.canvas_size[0] - 0.1) kps[..., 1] = np.clip(kps[..., 1], 0.1, self.canvas_size[1] - 0.1) # nbox, boxsize(4)*coord_dim(2) kps = kps.reshape(kps.shape[0], -1) # unflatten_seq(qkps) return kps, kp_labels, qkps, qkp_msks, def Polygonization(self, input_dict): ''' Process vertices. ''' vectors = input_dict['vectors'] n_lines = 0 for label, lines in vectors.items(): n_lines += len(lines) if not n_lines: input_dict['polys'] = [] return input_dict polyline_map, polyline_map_mask, polyline_map_weight = \ self.format_polyline_map(vectors) keypoint, keypoint_label, qkeypoint, qkeypoint_mask = \ self.format_keypoint(vectors) # gather polys = { # for det 'keypoint': keypoint, 'det_label': keypoint_label, # for gen 'gen_label': keypoint_label, 'qkeypoint': qkeypoint, 'qkeypoint_mask': qkeypoint_mask, 'polylines': polyline_map, # List[array] 'polyline_masks': polyline_map_mask, # List[array] 'polyline_weights': polyline_map_weight } # Format outputs input_dict['polys'] = polys return input_dict def __call__(self, input_dict): input_dict = self.Polygonization(input_dict) return input_dict def evaluate_line(polyline): edge = np.linalg.norm(polyline[1:] - polyline[:-1], axis=-1) start_end_weight = edge[(0, -1),].copy() mid_weight = (edge[:-1] + edge[1:]) * .5 pts_weight = np.concatenate( (start_end_weight[:1], mid_weight, start_end_weight[-1:])) denominator = pts_weight.sum() denominator = 1 if denominator == 0 else denominator pts_weight /= denominator # add weights for stop index pts_weight = np.repeat(pts_weight, 2) / 2 pts_weight = np.pad(pts_weight, ((0, 1)), constant_values=1 / (len(polyline) * 2)) return pts_weight def quantize_verts(verts, canvas_size, coord_dim): """Convert vertices from its original range ([-1,1]) to discrete values in [0, n_bits**2 - 1]. Args: verts (array): vertices coordinates, shape (seqlen, coords_dim) canvas_size (tuple): bev feature size coord_dim (int): dimension of point coordinates Returns: quantized_verts (array): quantized vertices, shape (seqlen, coords_dim) """ min_range = 0 max_range = 1 range_quantize = np.array(canvas_size) - 1 # (0-199) = 200 verts_ratio = (verts[:, :coord_dim] - min_range) / ( max_range - min_range) verts_quantize = verts_ratio * range_quantize[:coord_dim] return verts_quantize.astype('int32') def get_bbox(polyline, threshold): """Convert vertices from its original range ([-1,1]) to discrete values in [0, n_bits**2 - 1]. Args: polyline (array): point coordinates, shape (seqlen, 2) threshold (float): threshold for minimum bbox size Returns: bbox (array): bounding box in xyxy format, shape (2, 2) """ eps = 1e-4 polyline = LineString(polyline) bbox = polyline.bounds minx, miny, maxx, maxy = bbox W, H = maxx - minx, maxy - miny if W < threshold or H < threshold: remain = max((threshold - min(W, H)) / 2, eps) bbox = polyline.buffer(remain).envelope.bounds minx, miny, maxx, maxy = bbox bbox_np = np.array([[minx, miny], [maxx, maxy]]) bbox_np = np.clip(bbox_np, 0., 1.) return bbox_np