Enhance Shenron visualization pipeline with HD Matplotlib plots, unified 120m coordinate system, and FastHeatmapEngine UI caching for 2x packaging speed

scripts/test_shenron.py

@@ -102,6 +102,50 @@ def render_heatmap(data, vmin=None, vmax=None, cmap='viridis'):
     img.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("ascii")
 
+from matplotlib.figure import Figure
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+
+class FastHeatmapEngine:
+    """Stateful Matplotlib engine that reuses figure memory to achieve high-speed frame rendering."""
+    def __init__(self, extent, cmap='jet', vmin=None, vmax=None, title='Range-Azimuth Heatmap', xlabel='Lateral distance [m]', ylabel='Longitudinal distance [m]', xlim=None, ylim=None, aspect='auto', interpolation=None):
+        self.vmin = vmin
+        self.vmax = vmax
+        
+        self.fig = Figure(figsize=(6, 5), dpi=100)
+        self.canvas = FigureCanvasAgg(self.fig)
+        self.ax = self.fig.add_subplot(111)
+        
+        cm_obj = matplotlib.colormaps.get_cmap(cmap).copy()
+        cm_obj.set_bad(color='white')
+        
+        # Dummy matrix to initialize geometry
+        dummy = np.zeros((2, 2))
+        self.im = self.ax.imshow(dummy, extent=extent, cmap=cm_obj, vmin=vmin, vmax=vmax, origin='upper', aspect=aspect, interpolation=interpolation)
+        
+        if xlim is not None: self.ax.set_xlim(xlim)
+        if ylim is not None: self.ax.set_ylim(ylim)
+        
+        self.ax.set_xlabel(xlabel)
+        self.ax.set_ylabel(ylabel)
+        self.ax.set_title(title)
+        self.ax.grid(color='gray', linestyle='--', linewidth=0.5, alpha=0.7)
+        self.fig.colorbar(self.im, ax=self.ax, label='Magnitude (dB)')
+        self.fig.tight_layout()
+        
+    def render(self, data):
+        self.im.set_data(data)
+        if self.vmin is None and self.vmax is None:
+            v_low, v_high = np.nanmin(data), np.nanmax(data)
+            self.im.set_clim(v_low if not np.isnan(v_low) else 0.0, v_high if not np.isnan(v_high) else 1.0)
+            
+        self.fig.canvas.draw()
+        rgba = np.asarray(self.canvas.buffer_rgba())
+        img = Image.fromarray(rgba)
+        buf = io.BytesIO()
+        img.save(buf, format='png')
+        return base64.b64encode(buf.getvalue()).decode("ascii")
+
+
 def postprocess_ra(ra_heatmap, range_axis, smooth_sigma=1.0):
     """
     Refined RA post-processing pipeline for Physical Realism.
@@ -141,12 +185,14 @@ def postprocess_ra(ra_heatmap, range_axis, smooth_sigma=1.0):
     return ra_db
 
 
-def scan_convert_ra(ra_heatmap, range_axis, angle_axis, img_size=512):
+def scan_convert_ra(ra_heatmap, range_axis, angle_axis, img_size=512, max_display_range=None):
     """
     Polar-to-Cartesian scan conversion following FIG / Guide logic.
     Converts RA (Range, Angle) polar data into a 120° Fan-shaped Sector plot.
     """
-    max_range = range_axis[-1]
+    true_max_range = range_axis[-1]
+    max_range = max_display_range if max_display_range is not None else true_max_range
+    
     theta_min = angle_axis[0]
     theta_max = angle_axis[-1]
 
@@ -164,14 +210,14 @@ def scan_convert_ra(ra_heatmap, range_axis, angle_axis, img_size=512):
     fov_mask = (Theta_query >= theta_min) & (Theta_query <= theta_max) & (R_query <= max_range)
 
     # 5. Map RA Heatmap to Cartesian Grid
-    # Calculate fractional indices
-    r_idx = np.clip(((R_query / max_range) * (ra_heatmap.shape[0] - 1)).astype(int), 0, ra_heatmap.shape[0] - 1)
+    # Calculate fractional indices based on the true underlying data range
+    r_idx = np.clip(((R_query / true_max_range) * (ra_heatmap.shape[0] - 1)).astype(int), 0, ra_heatmap.shape[0] - 1)
     # theta index: Shift by theta_min to align 0..120 range
     theta_range = theta_max - theta_min
     theta_idx = np.clip(((Theta_query - theta_min) / theta_range * (ra_heatmap.shape[1] - 1)).astype(int), 0, ra_heatmap.shape[1] - 1)
 
     # Project
-    cartesian = np.full((img_size, img_size), np.min(ra_heatmap), dtype=np.float64)
+    cartesian = np.full((img_size, img_size), np.nan, dtype=np.float64)
     cartesian[fov_mask] = ra_heatmap[r_idx[fov_mask], theta_idx[fov_mask]]
 
     return cartesian
@@ -303,6 +349,7 @@ def run_testbench(iter_name):
             metrology_channels[r_type] = {
                 "rd": writer.register_channel(topic=f"/radar/{r_type}/heatmaps/range_doppler", message_encoding="json", schema_id=camera_schema_id),
                 "ra": writer.register_channel(topic=f"/radar/{r_type}/heatmaps/range_azimuth", message_encoding="json", schema_id=camera_schema_id),
+                "ra_dynamic": writer.register_channel(topic=f"/radar/{r_type}/heatmaps/range_azimuth_dynamic", message_encoding="json", schema_id=camera_schema_id),
                 "cfar": writer.register_channel(topic=f"/radar/{r_type}/heatmaps/cfar_mask", message_encoding="json", schema_id=camera_schema_id),
                 "telemetry": writer.register_channel(topic=f"/radar/{r_type}/metrics", message_encoding="json", schema_id=writer.register_schema(name="foxglove.Telemetry", encoding="jsonschema", data=json.dumps(FOXGLOVE_METRICS_SCHEMA).encode()))
             }
@@ -310,12 +357,15 @@ def run_testbench(iter_name):
         try:
             frames_gen = load_frames(logs_dir)
             frames = list(frames_gen)
+            if args.frames and args.frames > 0:
+                frames = frames[:args.frames]
         except Exception as e:
             print(f"[ERROR] Could not load frames.jsonl from {logs_dir}: {e}")
             return
             
-        # Pre-load physical axes (saved once per radar type during Stage 1)
+        # Pre-load physical axes & Heatmap Engines
         cached_axes = {}
+        render_engines = {}
         for r_type in radar_types:
             range_ax_p = iter_dir / r_type / "metrology" / "range_axis.npy"
             angle_ax_p = iter_dir / r_type / "metrology" / "angle_axis.npy"
@@ -327,6 +377,20 @@ def run_testbench(iter_name):
             else:
                 cached_axes[r_type] = None
                 
+            # Initialize the stateful Matplotlib renderers for extreme throughput
+            f_cfg = models[r_type].radar_obj.f if r_type in models else 77e9
+            chirp_rep_cfg = models[r_type].radar_obj.chirp_rep if r_type in models else 3e-5
+            max_vel_cfg = (3e8 / f_cfg) / (4 * chirp_rep_cfg)
+            max_r_cfg = cached_axes[r_type]['range_axis'][-1] if cached_axes[r_type] else 150
+            display_limit_cfg = 120.0
+            
+            render_engines[r_type] = {
+                'rd': FastHeatmapEngine(extent=[-max_vel_cfg, max_vel_cfg, 0, max_r_cfg], cmap='viridis', title=f'{r_type.upper()} Range-Doppler', xlabel='Doppler Velocity [m/s]', ylabel='Range [m]', ylim=[0, 120], interpolation='bicubic'),
+                'cfar': FastHeatmapEngine(extent=[-max_vel_cfg, max_vel_cfg, 0, max_r_cfg], cmap='plasma', title=f'{r_type.upper()} CFAR Noise Threshold', xlabel='Doppler Velocity [m/s]', ylabel='Range [m]', ylim=[0, 120], interpolation='bicubic'),
+                'ra_static': FastHeatmapEngine(extent=[-display_limit_cfg, display_limit_cfg, 0, display_limit_cfg], cmap='jet', vmin=-5, vmax=45, title=f'{r_type.upper()} Range-Azimuth (Absolute)', xlabel='Lateral distance [m]', ylabel='Longitudinal distance [m]', aspect='equal'),
+                'ra_dyn': FastHeatmapEngine(extent=[-display_limit_cfg, display_limit_cfg, 0, display_limit_cfg], cmap='jet', title=f'{r_type.upper()} Range-Azimuth (Dynamic)', xlabel='Lateral distance [m]', ylabel='Longitudinal distance [m]', aspect='equal')
+            }
+
         for frame in tqdm.tqdm(frames, desc="  Packaging Frames", unit="frame"):
             ts_ns = int(frame["timestamp"] * 1e9)
             ts_sec = ts_ns // 1_000_000_000
@@ -442,10 +506,15 @@ def run_testbench(iter_name):
                     ra_p = met_folder / "ra" / shenron_fname
                     cf_p = met_folder / "cfar" / shenron_fname
                     
+                    
                     if rd_p.exists():
                         rd_data = np.load(rd_p)
-                        # Flip UD so Range 0 (ego) is at the bottom
-                        b64 = render_heatmap(np.log10(np.flipud(rd_data) + 1e-9), cmap='viridis')
+                        # Apply log conversion minus identical system gain offset to maintain -5 to 45 scaling
+                        # Simple 10*log10 - SYSTEM_GAIN_OFFSET
+                        rd_db = 10 * np.log10(np.clip(rd_data, 1e-9, None)) - 68.0 
+                        
+                        # Flip UD so Range 0 (ego) is at the bottom. Use Cached Renderer.
+                        b64 = render_engines[r_type]['rd'].render(np.flipud(rd_db))
                         if b64:
                             msg = {"timestamp": {"sec": ts_sec, "nsec": ts_nsec}, "frame_id": "ego_vehicle", "format": "png", "data": b64}
                             writer.add_message(metrology_channels[r_type]["rd"], log_time=ts_ns, data=json.dumps(msg).encode(), publish_time=ts_ns)
@@ -458,8 +527,22 @@ def run_testbench(iter_name):
                             # Apply full post-processing chain (log, R² compensation, clutter, normalize, smooth)
                             ra_processed = postprocess_ra(ra_data, axes['range_axis'], smooth_sigma=0.0) # Disabled smoothing as per focus fix
                             # Polar Sector BEV plot — geometrically accurate
-                            bev_data = scan_convert_ra(ra_processed, axes['range_axis'], axes['angle_axis'], img_size=512)
-                            b64 = render_heatmap(bev_data, cmap='jet', vmin=-5, vmax=45)
+                            # Project using 512x512 resolution constrained entirely to the 120m boundary to avoid pixelation
+                            display_rng_limit = 120.0
+                            bev_data = scan_convert_ra(ra_processed, axes['range_axis'], axes['angle_axis'], img_size=512, max_display_range=display_rng_limit)
+                            
+                            # 1. PRIMARY PLOT: Static fixed bounds for 1:1 magnitude tracking over time
+                            b64_static = render_engines[r_type]['ra_static'].render(bev_data)
+                            if b64_static:
+                                msg = {"timestamp": {"sec": ts_sec, "nsec": ts_nsec}, "frame_id": "ego_vehicle", "format": "png", "data": b64_static}
+                                writer.add_message(metrology_channels[r_type]["ra"], log_time=ts_ns, data=json.dumps(msg).encode(), publish_time=ts_ns)
+                            
+                            # 2. HIGHLIGHT PLOT: Dynamic bounds to track the peak signature without external thresholds
+                            b64_dynamic = render_engines[r_type]['ra_dyn'].render(bev_data)
+                            if b64_dynamic:
+                                msg_dyn = {"timestamp": {"sec": ts_sec, "nsec": ts_nsec}, "frame_id": "ego_vehicle", "format": "png", "data": b64_dynamic}
+                                writer.add_message(metrology_channels[r_type]["ra_dynamic"], log_time=ts_ns, data=json.dumps(msg_dyn).encode(), publish_time=ts_ns)
+
                         else:
                             # Fallback: rectangular log plot (no axis info available)
                             b64 = render_heatmap(np.log10(np.flipud(ra_data) + 1e-9), cmap='magma')
@@ -469,8 +552,13 @@ def run_testbench(iter_name):
                     
                     if cf_p.exists():
                         cf_data = np.load(cf_p)
+                        
+                        # Convert threshold power floor to pure threshold DB mask similar to RD
+                        cf_db = 10 * np.log10(np.clip(cf_data, 1e-9, None)) - 68.0 
+                        
                         # Flip UD so Range 0 (ego) is at the bottom
-                        b64 = render_heatmap(np.log10(np.flipud(cf_data) + 1e-9), cmap='plasma') # Plasma for threshold mask
+                        # Revert to dynamic (None) scaling so threshold logic is easily visible
+                        b64 = render_engines[r_type]['cfar'].render(np.flipud(cf_db))
                         if b64:
                             msg = {"timestamp": {"sec": ts_sec, "nsec": ts_nsec}, "frame_id": "ego_vehicle", "format": "png", "data": b64}
                             writer.add_message(metrology_channels[r_type]["cfar"], log_time=ts_ns, data=json.dumps(msg).encode(), publish_time=ts_ns)