| | |
| |
|
| | 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) |
| |
|
