examples/a2/algo/heatmap/heatmap_visualization.py

examples/a2/algo/heatmap/heatmap_visualization.py#

  1# Copyright (c) Acconeer AB, 2024-2026
  2# All rights reserved
  3
  4from __future__ import annotations
  5
  6import argparse
  7import os
  8import time
  9
 10import numpy as np
 11
 12from PySide6.QtGui import QTransform
 13
 14import pyqtgraph as pg
 15
 16import acconeer.exptool as et
 17from acconeer.exptool import a2
 18from acconeer.exptool.a2.algo.heatmap.conventional import (
 19    HeatmapConfig,
 20    HeatmapProcessor,
 21    ProcessorResult,
 22)
 23
 24
 25""" Heatmap Config to use for heatmap generation."""
 26HEATMAP_CONFIG = HeatmapConfig(
 27    antenna_array=a2.AntennaArray.AP212,
 28    heatmap_mode=HeatmapConfig.HeatmapMode.MOVEMENT,
 29    covariance_matrix_time_constant_s=0.15,
 30    high_pass_filter_time_constant_s=1.0,
 31    az_angle_range=[-60.0, 60.0],
 32    el_angle_range=[0.0, 0.0],
 33    num_az_angles=25,
 34    num_el_angles=1,
 35)
 36
 37DEGREE_SYMBOL = "\u00b0"
 38BASE_STEP_LENGTH = 2.5e-3
 39SKIN = os.environ.get("ACC_SKIN", "none").lower()
 40TEXT_COLOR = "white" if SKIN == "tron" else "black"
 41
 42
 43class PGUpdater:
 44    """
 45    Minimal version of plotting from exploration tool.
 46    """
 47
 48    def __init__(self, session_config, heatmap_config) -> None:
 49        self._config: a2.SessionConfig | None = session_config
 50        self.heatmap_config = heatmap_config
 51
 52    def setup(self, win):
 53        self.viridis_heatmap_image_item = pg.ImageItem()
 54        self.viridis_heatmap_image_item.setLookupTable(et.utils.pg_mpl_cmap("viridis"))
 55
 56        self.viridis_heatmap_plot_item = win.addPlot(row=0, col=0, title="Heatmap Visualization")
 57        self.viridis_heatmap_plot_item.getAxis("left").setLabel("Distance", units="m")
 58        self.viridis_heatmap_plot_item.getAxis("bottom").setLabel("Az angle", units=DEGREE_SYMBOL)
 59        self.viridis_heatmap_colorbar = self.viridis_heatmap_plot_item.addColorBar(
 60            self.viridis_heatmap_image_item, colorMap="viridis"
 61        )
 62
 63        self.viridis_heatmap_plot_item.clear()
 64        self.viridis_heatmap_plot_item.addItem(self.viridis_heatmap_image_item)
 65        self.viridis_heatmap_plot_smooth_max = et.utils.SmoothMax(
 66            self._config.update_rate, tau_decay=0.1, tau_grow=0.1
 67        )
 68
 69        start_point = self._config.sensor_config.subsweep.start_point
 70        step_length = self._config.sensor_config.subsweep.step_length
 71
 72        self.viridis_heatmap_image_item.setTransform(
 73            QTransform()
 74            .translate(-self.heatmap_config.az_span / 2, start_point * BASE_STEP_LENGTH)
 75            .scale(
 76                self.heatmap_config.az_span / len(self.heatmap_config.az_angles),
 77                step_length * BASE_STEP_LENGTH,
 78            )
 79        )
 80
 81    def update(self, processor_result: ProcessorResult) -> None:
 82        """Plotting is achieved by repeated calls to update with new processorresults"""
 83
 84        if processor_result.heatmap is None:
 85            return
 86        heatmap = processor_result.heatmap[:, 0, :]
 87        max_val = np.max(heatmap)
 88
 89        max_heat = self.viridis_heatmap_plot_smooth_max.update(max_val)
 90        image = np.nan_to_num(heatmap)
 91        self.viridis_heatmap_colorbar.setLevels((0, max_heat))
 92        self.viridis_heatmap_image_item.updateImage(image, levels=(0, max_heat))
 93
 94
 95def main():
 96    parser = argparse.ArgumentParser()
 97    parser.add_argument(
 98        "--input-file",
 99        required=True,
100        help="path to h5 record with data to visualize with heatmap processor.",
101    )
102
103    args = parser.parse_args()
104
105    input_file = args.input_file
106
107    # Load the recorded data
108    record = a2.open_record(input_file)
109
110    # Figure out update rate
111    ticks_per_second = record.server_info.ticks_per_second
112    session_index = 0  # Can be different but is usually zero
113    tick_period = record.session(session_index).tick_period
114    if tick_period > 0.0:
115        update_rate = ticks_per_second / tick_period
116    else:
117        update_rate = 10  # Default value, manually adjust here if required.
118
119    session_config = record.session_config
120    subsweep_config = (
121        session_config.sensor_config.subsweep
122    )  # Assume that we only have a single subsweep.
123    processor = HeatmapProcessor(
124        subsweep_config=subsweep_config,
125        heatmap_config=HEATMAP_CONFIG,
126        update_rate=update_rate,
127    )
128
129    # Set up the heatmap plot
130    pg_updater = PGUpdater(session_config, HEATMAP_CONFIG)
131    pg_process = et.PGProcess(pg_updater)
132    pg_process.start()
133
134    interrupt_handler = et.utils.ExampleInterruptHandler()
135    print("Press Ctrl-C to end session")
136
137    while not interrupt_handler.got_signal:
138        # Go through all of the recorded data
139        for result in record.results:
140            # Post-process radar data and get the heatmap result
141            processor_result = processor.process(result)
142
143            # Make sure we have a heatmap to plot
144            if processor_result.heatmap is None:
145                continue
146
147            try:
148                # Plot the heatmap
149                pg_process.put_data(processor_result)
150            except et.PGProccessDiedException:
151                break
152            time.sleep(1.0 / update_rate)
153
154    print("Disconnecting...")
155    pg_process.close()
156
157
158if __name__ == "__main__":
159    main()

View this example on GitHub: acconeer/acconeer-python-exploration