| 6 | |
| 7 | |
| 8 | class MapEncoder(nn.Module): |
| 9 | def __init__( |
| 10 | self, |
| 11 | polygon_channel=6, |
| 12 | dim=128, |
| 13 | use_lane_boundary=False, |
| 14 | ) -> None: |
| 15 | super().__init__() |
| 16 | |
| 17 | self.dim = dim |
| 18 | self.use_lane_boundary = use_lane_boundary |
| 19 | self.polygon_channel = ( |
| 20 | polygon_channel + 4 if use_lane_boundary else polygon_channel |
| 21 | ) |
| 22 | |
| 23 | self.polygon_encoder = PointsEncoder(self.polygon_channel, dim) |
| 24 | self.speed_limit_emb = FourierEmbedding(1, dim, 64) |
| 25 | |
| 26 | self.type_emb = nn.Embedding(3, dim) |
| 27 | self.on_route_emb = nn.Embedding(2, dim) |
| 28 | self.traffic_light_emb = nn.Embedding(4, dim) |
| 29 | self.unknown_speed_emb = nn.Embedding(1, dim) |
| 30 | |
| 31 | def forward(self, data) -> torch.Tensor: |
| 32 | polygon_center = data["map"]["polygon_center"] |
| 33 | polygon_type = data["map"]["polygon_type"].long() |
| 34 | polygon_on_route = data["map"]["polygon_on_route"].long() |
| 35 | polygon_tl_status = data["map"]["polygon_tl_status"].long() |
| 36 | polygon_has_speed_limit = data["map"]["polygon_has_speed_limit"] |
| 37 | polygon_speed_limit = data["map"]["polygon_speed_limit"] |
| 38 | point_position = data["map"]["point_position"] |
| 39 | point_vector = data["map"]["point_vector"] |
| 40 | point_orientation = data["map"]["point_orientation"] |
| 41 | valid_mask = data["map"]["valid_mask"] |
| 42 | |
| 43 | if self.use_lane_boundary: |
| 44 | polygon_feature = torch.cat( |
| 45 | [ |
| 46 | point_position[:, :, 0] - polygon_center[..., None, :2], |
| 47 | point_vector[:, :, 0], |
| 48 | torch.stack( |
| 49 | [ |
| 50 | point_orientation[:, :, 0].cos(), |
| 51 | point_orientation[:, :, 0].sin(), |
| 52 | ], |
| 53 | dim=-1, |
| 54 | ), |
| 55 | point_position[:, :, 1] - point_position[:, :, 0], |
| 56 | point_position[:, :, 2] - point_position[:, :, 0], |
| 57 | ], |
| 58 | dim=-1, |
| 59 | ) |
| 60 | else: |
| 61 | polygon_feature = torch.cat( |
| 62 | [ |
| 63 | point_position[:, :, 0] - polygon_center[..., None, :2], |
| 64 | point_vector[:, :, 0], |
| 65 | torch.stack( |