diff --git a/scripts/test_shenron.py b/scripts/test_shenron.py index b6abf3b..f623219 100644 --- a/scripts/test_shenron.py +++ b/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" @@ -326,6 +376,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) @@ -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,19 +527,38 @@ 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') - 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]["ra"], log_time=ts_ns, data=json.dumps(msg).encode(), publish_time=ts_ns) + 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]["ra"], log_time=ts_ns, data=json.dumps(msg).encode(), publish_time=ts_ns) 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)