Class VisualizerTrajectory

flowmap/visualization/visualizer_trajectory.py:59–109

class VisualizerTrajectory(Visualizer[VisualizerTrajectoryCfg]):
    def __init__(self, cfg: VisualizerTrajectoryCfg) -> None:
        super().__init__(cfg)
        self.ates = []

    def visualize(
        self,
        batch: Batch,
        flows: Flows,
        tracks: list[Tracks] | None,
        model_output: ModelOutput,
        model: Model,
        global_step: int,
    ) -> dict[str, Float[Tensor, "3 _ _"] | Float[Tensor, ""]]:
        # If there's no ground truth, do nothing.
        if batch.extrinsics is None:
            return {}

        # For now, only support batch size 1 for visualization.
        b, _, _, _, _ = batch.videos.shape
        assert b == 1

        # Compute the ATE.
        try:
            ate, positions_gt, positions_hat = compute_ate(
                batch.extrinsics[0, :, :3, 3],
                model_output.extrinsics[0, :, :3, 3],
            )
        except ValueError:
            return {}
        result = {"metrics/ate": ate}

        # Visualize the trajectory.
        if self.cfg.generate_plot:
            fg = generate_plot(
                [positions_gt.cpu().numpy(), positions_hat.cpu().numpy()],
                ["Ground-truth", "Predicted"],
            )
            visualization = fig_to_image(fg)
            plt.close(fg)
            result["trajectory"] = add_border(visualization)

        if self.cfg.ate_save_path is not None:
            # It's not ideal to write out a file during each optimization step, but this
            # only needs to be run once to generate a plot in the paper, so it's fine.
            self.ates.append(ate.item())
            self.cfg.ate_save_path.parent.mkdir(exist_ok=True, parents=True)
            with self.cfg.ate_save_path.open("w") as f:
                json.dump(self.ates, f)

        return result