|
|
import torch
|
|
|
|
|
|
from .metrics import PixelAccuracy, MeanObservableIOU, MeanUnobservableIOU, ObservableIOU, UnobservableIOU, mAP
|
|
|
|
|
|
from .loss import EnhancedLoss
|
|
|
|
|
|
from .segmentation_head import SegmentationHead
|
|
|
|
|
|
from . import get_model
|
|
|
from .base import BaseModel
|
|
|
from .bev_projection import CartesianProjection, PolarProjectionDepth
|
|
|
from .schema import ModelConfiguration
|
|
|
|
|
|
class MapPerceptionNet(BaseModel):
|
|
|
|
|
|
def _init(self, conf: ModelConfiguration):
|
|
|
self.image_encoder = get_model(
|
|
|
conf.image_encoder.name
|
|
|
)(conf.image_encoder.backbone)
|
|
|
|
|
|
self.decoder = SegmentationHead(
|
|
|
in_channels=conf.latent_dim, n_classes=conf.num_classes)
|
|
|
|
|
|
ppm = conf.pixel_per_meter
|
|
|
self.projection_polar = PolarProjectionDepth(
|
|
|
conf.z_max,
|
|
|
ppm,
|
|
|
conf.scale_range,
|
|
|
conf.z_min,
|
|
|
)
|
|
|
self.projection_bev = CartesianProjection(
|
|
|
conf.z_max, conf.x_max, ppm, conf.z_min
|
|
|
)
|
|
|
|
|
|
self.scale_classifier = torch.nn.Linear(
|
|
|
conf.latent_dim, conf.num_scale_bins
|
|
|
)
|
|
|
|
|
|
self.num_classes = conf.num_classes
|
|
|
|
|
|
self.loss_fn = EnhancedLoss(conf.loss)
|
|
|
|
|
|
def _forward(self, data):
|
|
|
f_image, camera = self.image_encoder(data)
|
|
|
|
|
|
scales = self.scale_classifier(
|
|
|
f_image.moveaxis(1, -1))
|
|
|
f_polar = self.projection_polar(f_image, scales, camera)
|
|
|
|
|
|
|
|
|
f_bev, valid_bev, _ = self.projection_bev(
|
|
|
f_polar.float(), None, camera.float()
|
|
|
)
|
|
|
|
|
|
output = self.decoder(f_bev[..., :-1])
|
|
|
|
|
|
probs = torch.nn.functional.sigmoid(output)
|
|
|
|
|
|
return {
|
|
|
"output": probs,
|
|
|
"logits": output,
|
|
|
"scales": scales,
|
|
|
"features_image": f_image,
|
|
|
"features_bev": f_bev,
|
|
|
"valid_bev": valid_bev.squeeze(1),
|
|
|
}
|
|
|
|
|
|
def loss(self, pred, data):
|
|
|
loss = self.loss_fn(pred, data)
|
|
|
return loss
|
|
|
|
|
|
def metrics(self):
|
|
|
m = {
|
|
|
"pix_acc": PixelAccuracy(),
|
|
|
"map": mAP(self.num_classes),
|
|
|
"miou_observable": MeanObservableIOU(self.num_classes),
|
|
|
"miou_non_observable": MeanUnobservableIOU(self.num_classes),
|
|
|
}
|
|
|
m.update(
|
|
|
{
|
|
|
f"IoU_observable_class_{i}": ObservableIOU(i, num_classes=self.num_classes)
|
|
|
for i in range(self.num_classes)
|
|
|
}
|
|
|
)
|
|
|
m.update(
|
|
|
{
|
|
|
f"IoU_non_observable_{i}": UnobservableIOU(i, num_classes=self.num_classes)
|
|
|
for i in range(self.num_classes)
|
|
|
}
|
|
|
)
|
|
|
return m
|
|
|
|
