| |
|
|
| import torch |
| import torchmetrics |
| from torchmetrics.utilities.data import dim_zero_cat |
|
|
| from .utils import deg2rad, rotmat2d |
|
|
|
|
| def location_error(uv, uv_gt, ppm=1): |
| return torch.norm(uv - uv_gt.to(uv), dim=-1) / ppm |
|
|
| def location_error_single(uv, uv_gt, ppm=1): |
| return torch.norm(uv - uv_gt.to(uv), dim=-1) / ppm |
|
|
| def angle_error(t, t_gt): |
| error = torch.abs(t % 360 - t_gt.to(t) % 360) |
| error = torch.minimum(error, 360 - error) |
| return error |
|
|
|
|
| class Location2DRecall(torchmetrics.MeanMetric): |
| def __init__(self, threshold, pixel_per_meter, key="uv_max", *args, **kwargs): |
| self.threshold = threshold |
| self.ppm = pixel_per_meter |
| self.key = key |
| super().__init__(*args, **kwargs) |
|
|
| def update(self, pred, data): |
| self.cuda() |
| error = location_error(pred[self.key], data["uv"], self.ppm) |
| |
| super().update((error <= torch.tensor(self.threshold,device=error.device)).float()) |
|
|
| class Location1DRecall(torchmetrics.MeanMetric): |
| def __init__(self, threshold, pixel_per_meter, key="uv_max", *args, **kwargs): |
| self.threshold = threshold |
| self.ppm = pixel_per_meter |
| self.key = key |
| super().__init__(*args, **kwargs) |
|
|
| def update(self, pred, data): |
| self.cuda() |
| error = location_error(pred[self.key], data["uv"], self.ppm) |
| |
| super().update((error <= torch.tensor(self.threshold,device=error.device)).float()) |
| class AngleRecall(torchmetrics.MeanMetric): |
| def __init__(self, threshold, key="yaw_max", *args, **kwargs): |
| self.threshold = threshold |
| self.key = key |
|
|
| super().__init__(*args, **kwargs) |
|
|
| def update(self, pred, data): |
| self.cuda() |
| error = angle_error(pred[self.key], data["roll_pitch_yaw"][..., -1]) |
| super().update((error <= self.threshold).float()) |
|
|
|
|
| class MeanMetricWithRecall(torchmetrics.Metric): |
| full_state_update = True |
|
|
| def __init__(self): |
| super().__init__() |
| self.add_state("value", default=[], dist_reduce_fx="cat") |
| def compute(self): |
| return dim_zero_cat(self.value).mean(0) |
|
|
| def get_errors(self): |
| return dim_zero_cat(self.value) |
|
|
| def recall(self, thresholds): |
| self.cuda() |
| error = self.get_errors() |
| thresholds = error.new_tensor(thresholds) |
| return (error.unsqueeze(-1) < thresholds).float().mean(0) * 100 |
|
|
|
|
| class AngleError(MeanMetricWithRecall): |
| def __init__(self, key): |
| super().__init__() |
| self.key = key |
|
|
| def update(self, pred, data): |
| self.cuda() |
| value = angle_error(pred[self.key], data["roll_pitch_yaw"][..., -1]) |
| if value.numel(): |
| self.value.append(value) |
|
|
|
|
| class Location2DError(MeanMetricWithRecall): |
| def __init__(self, key, pixel_per_meter): |
| super().__init__() |
| self.key = key |
| self.ppm = pixel_per_meter |
|
|
| def update(self, pred, data): |
| self.cuda() |
| value = location_error(pred[self.key], data["uv"], self.ppm) |
| if value.numel(): |
| self.value.append(value) |
|
|
|
|
| class LateralLongitudinalError(MeanMetricWithRecall): |
| def __init__(self, pixel_per_meter, key="uv_max"): |
| super().__init__() |
| self.ppm = pixel_per_meter |
| self.key = key |
|
|
| def update(self, pred, data): |
| self.cuda() |
| yaw = deg2rad(data["roll_pitch_yaw"][..., -1]) |
| shift = (pred[self.key] - data["uv"]) * yaw.new_tensor([-1, 1]) |
| shift = (rotmat2d(yaw) @ shift.unsqueeze(-1)).squeeze(-1) |
| error = torch.abs(shift) / self.ppm |
| value = error.view(-1, 2) |
| if value.numel(): |
| self.value.append(value) |
|
|
