| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| from .cluster import LaneNetPostProcessor |
| from .connect import sort_points_by_dist, connect_by_direction |
|
|
|
|
| def onehot_encoding(logits, dim=0): |
| max_idx = torch.argmax(logits, dim, keepdim=True) |
| one_hot = logits.new_full(logits.shape, 0) |
| one_hot.scatter_(dim, max_idx, 1) |
| return one_hot |
|
|
|
|
| def onehot_encoding_spread(logits, dim=1): |
| max_idx = torch.argmax(logits, dim, keepdim=True) |
| one_hot = logits.new_full(logits.shape, 0) |
| one_hot.scatter_(dim, max_idx, 1) |
| one_hot.scatter_(dim, torch.clamp(max_idx-1, min=0), 1) |
| one_hot.scatter_(dim, torch.clamp(max_idx-2, min=0), 1) |
| one_hot.scatter_(dim, torch.clamp(max_idx+1, max=logits.shape[dim]-1), 1) |
| one_hot.scatter_(dim, torch.clamp(max_idx+2, max=logits.shape[dim]-1), 1) |
|
|
| return one_hot |
|
|
|
|
| def get_pred_top2_direction(direction, dim=1): |
| direction = torch.softmax(direction, dim) |
| idx1 = torch.argmax(direction, dim) |
| idx1_onehot_spread = onehot_encoding_spread(direction, dim) |
| idx1_onehot_spread = idx1_onehot_spread.bool() |
| direction[idx1_onehot_spread] = 0 |
| idx2 = torch.argmax(direction, dim) |
| direction = torch.stack([idx1, idx2], dim) - 1 |
| return direction |
|
|
|
|
| def vectorize(segmentation, embedding, direction, angle_class): |
| segmentation = segmentation.softmax(0) |
| embedding = embedding.cpu() |
| direction = direction.permute(1, 2, 0).cpu() |
| direction = get_pred_top2_direction(direction, dim=-1) |
|
|
| max_pool_1 = nn.MaxPool2d((1, 5), padding=(0, 2), stride=1) |
| avg_pool_1 = nn.AvgPool2d((9, 5), padding=(4, 2), stride=1) |
| max_pool_2 = nn.MaxPool2d((5, 1), padding=(2, 0), stride=1) |
| avg_pool_2 = nn.AvgPool2d((5, 9), padding=(2, 4), stride=1) |
| post_processor = LaneNetPostProcessor(dbscan_eps=1.5, postprocess_min_samples=50) |
|
|
| oh_pred = onehot_encoding(segmentation).cpu().numpy() |
| confidences = [] |
| line_types = [] |
| simplified_coords = [] |
| for i in range(1, oh_pred.shape[0]): |
| single_mask = oh_pred[i].astype('uint8') |
| single_embedding = embedding.permute(1, 2, 0) |
|
|
| single_class_inst_mask, single_class_inst_coords = post_processor.postprocess(single_mask, single_embedding) |
| if single_class_inst_mask is None: |
| continue |
|
|
| num_inst = len(single_class_inst_coords) |
|
|
| prob = segmentation[i] |
| prob[single_class_inst_mask == 0] = 0 |
| nms_mask_1 = ((max_pool_1(prob.unsqueeze(0))[0] - prob) < 0.0001).cpu().numpy() |
| avg_mask_1 = avg_pool_1(prob.unsqueeze(0))[0].cpu().numpy() |
| nms_mask_2 = ((max_pool_2(prob.unsqueeze(0))[0] - prob) < 0.0001).cpu().numpy() |
| avg_mask_2 = avg_pool_2(prob.unsqueeze(0))[0].cpu().numpy() |
| vertical_mask = avg_mask_1 > avg_mask_2 |
| horizontal_mask = ~vertical_mask |
| nms_mask = (vertical_mask & nms_mask_1) | (horizontal_mask & nms_mask_2) |
|
|
| for j in range(1, num_inst + 1): |
| full_idx = np.where((single_class_inst_mask == j)) |
| full_lane_coord = np.vstack((full_idx[1], full_idx[0])).transpose() |
| confidence = prob[single_class_inst_mask == j].mean().item() |
|
|
| idx = np.where(nms_mask & (single_class_inst_mask == j)) |
| if len(idx[0]) == 0: |
| continue |
| lane_coordinate = np.vstack((idx[1], idx[0])).transpose() |
|
|
| range_0 = np.max(full_lane_coord[:, 0]) - np.min(full_lane_coord[:, 0]) |
| range_1 = np.max(full_lane_coord[:, 1]) - np.min(full_lane_coord[:, 1]) |
| if range_0 > range_1: |
| lane_coordinate = sorted(lane_coordinate, key=lambda x: x[0]) |
| else: |
| lane_coordinate = sorted(lane_coordinate, key=lambda x: x[1]) |
|
|
| lane_coordinate = np.stack(lane_coordinate) |
| lane_coordinate = sort_points_by_dist(lane_coordinate) |
| lane_coordinate = lane_coordinate.astype('int32') |
| lane_coordinate = connect_by_direction(lane_coordinate, direction, step=7, per_deg=360 / angle_class) |
|
|
| simplified_coords.append(lane_coordinate) |
| confidences.append(confidence) |
| line_types.append(i-1) |
|
|
| return simplified_coords, confidences, line_types |
|
|
