""" create_dataset.py Unified dataset creation script for the UAV acoustic localization dataset. Usage: python create_dataset.py --config [options] # Process a single session using its alignment_params.json: python create_dataset.py --config dataset/test_1/Oct_11_2024/flight_5/alignment_params.json # Process a specific segment (for sessions with multiple segments like no-drone): python create_dataset.py --config dataset/test_2/Mar_18_2025/flight_2/alignment_params.json --segment no_drone_1 # Override output folder: python create_dataset.py --config dataset/test_1/Oct_11_2024/flight_5/alignment_params.json --output /path/to/npz/ # Run without generating the visualization MP4 (faster): python create_dataset.py --config ... --no-video """ import os import json import argparse import numpy as np import matplotlib.pyplot as plt import pandas as pd from pyproj import Transformer from matplotlib.animation import FFMpegWriter from numba import njit from audio_beamforming import beamform_time from geo_utils import ( wrap_angle, calculate_angle_difference, calculate_azimuth_meters, calculate_elevation_meters, calculate_total_distance_meters, ) from io_utils import ( read_wav_block, apply_bandpass_filter, calculate_time, initialize_beamforming_params, open_wav_files, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def skip_wav_seconds(wav_file, seconds, rate, extra_samples=0): frames_to_skip = int(seconds * rate) + extra_samples wav_file.setpos(frames_to_skip) def format_time_s(total_seconds: float) -> str: minutes = int(total_seconds // 60) seconds = total_seconds % 60 return f"{minutes:02d}:{seconds:05.2f}" def calculate_angular_distance(az1, el1, az2, el2): az1_r, el1_r = np.radians(az1), np.radians(el1) az2_r, el2_r = np.radians(az2), np.radians(el2) x1, y1, z1 = np.cos(el1_r)*np.cos(az1_r), np.cos(el1_r)*np.sin(az1_r), np.sin(el1_r) x2, y2, z2 = np.cos(el2_r)*np.cos(az2_r), np.cos(el2_r)*np.sin(az2_r), np.sin(el2_r) dot = np.clip(x1*x2 + y1*y2 + z1*z2, -1.0, 1.0) return np.degrees(np.arccos(dot)) @njit def shift_signal_beamforming(signal, delay_samples): num_samples = signal.shape[0] shifted = np.zeros_like(signal) if delay_samples > 0: if delay_samples < num_samples: shifted[delay_samples:] = signal[:-delay_samples] elif delay_samples < 0: ds = -delay_samples if ds < num_samples: shifted[:-ds] = signal[ds:] else: for i in range(num_samples): shifted[i] = signal[i] return shifted def calculate_delays_for_direction(mic_positions, azimuth, elevation, sample_rate, speed_of_sound): az_r, el_r = np.radians(azimuth), np.radians(elevation) direction = np.array([np.cos(el_r)*np.cos(az_r), np.cos(el_r)*np.sin(az_r), np.sin(el_r)]) delays = np.dot(mic_positions, direction) / speed_of_sound return np.round(delays * sample_rate).astype(np.int32) def apply_beamforming(signal_data, delay_samples): num_samples, num_mics = signal_data.shape output = np.zeros(num_samples, dtype=np.float64) for mic_idx in range(num_mics): output += shift_signal_beamforming(signal_data[:, mic_idx], delay_samples[mic_idx]) return output / num_mics def beamform_in_direction(signal_data, mic_positions, azimuth, elevation, sample_rate, speed_of_sound): delays = calculate_delays_for_direction(mic_positions, azimuth, elevation, sample_rate, speed_of_sound) return apply_beamforming(signal_data, delays).astype(np.float32) # --------------------------------------------------------------------------- # Config loading # --------------------------------------------------------------------------- def load_segment_config(params_path: str, segment_name: str | None) -> dict: """ Load and resolve a session config from alignment_params.json. Returns a flat drone_config dict ready for process_session(). Supports both flat configs (single segment) and configs with a 'segments' list. """ config_dir = os.path.dirname(os.path.abspath(params_path)) with open(params_path) as f: params = json.load(f) def resolve(path): """Make a path absolute relative to the config file location.""" if path is None: return None if os.path.isabs(path): return path return os.path.normpath(os.path.join(config_dir, path)) base = { "sample_rate_hz": params.get("sample_rate_hz", 48000), "corrections_samples": params.get("corrections_samples", [0, 0, 0, 0]), "wavs_nosync": [resolve(w) for w in params.get("wavs_nosync", [])], "ref_csv": resolve(params.get("ref_csv")), "session": params.get("session", ""), "mic_array": params.get("mic_array"), } # Config has a 'segments' list (multi-segment, e.g. drone + no-drone in same WAVs) if "segments" in params: available = {s["name"]: s for s in params["segments"]} if segment_name is None: names = list(available.keys()) if len(names) == 1: segment_name = names[0] else: raise ValueError( f"Config has multiple segments: {names}. " f"Specify one with --segment." ) if segment_name not in available: raise ValueError(f"Segment '{segment_name}' not found in {params_path}. " f"Available: {list(available.keys())}") seg = available[segment_name] base.update({ "segment_name": segment_name, "skip_seconds": seg["skip_seconds"], "start_index_csv": seg.get("start_index_csv", 0), "initial_azimuth": seg.get("initial_azimuth_deg", 0.0), "initial_elevation": seg.get("initial_elevation_deg", 0.0), "flight_csv": resolve(seg.get("flight_csv")), "use_csv_log": seg.get("flight_csv") is not None, }) else: # Flat config (single segment) base.update({ "segment_name": params.get("flight_id", "flight"), "skip_seconds": params.get("skip_seconds", 0.0), "start_index_csv": params.get("start_index_csv", 0), "initial_azimuth": params.get("initial_azimuth_deg", 0.0), "initial_elevation": params.get("initial_elevation_deg", 0.0), "flight_csv": resolve(params.get("flight_csv")), "use_csv_log": params.get("flight_csv") is not None, }) return base # --------------------------------------------------------------------------- # Core processing # --------------------------------------------------------------------------- EXPERIMENT_PARAMS = { "CHANNELS": 6, "sample_rate": 48000, "chunk_duration_s": 0.1, "speed_of_sound": 343, "RECORD_SECONDS": 1_200_000, "lowcut": 200.0, "highcut": 8000.0, "azimuth_range": np.arange(-180, 181, 4), "elevation_range": np.arange(0, 91, 4), "angular_threshold": 10, } # Fallback mic array geometry, used only if a session's alignment_params.json has no # "mic_array" block. The array was recalibrated between sessions — see "mic_array" in # each alignment_params.json for the actual per-session values used. DEFAULT_MIC_ARRAY = dict( a_=[0, -120, -240], a2_=[-40, -80, -160, -200, -280, -320], h_=[1.12, 1.02, 0.87, 0.68, 0.47, 0.02], r_=[0.1, 0.16, 0.23, 0.29, 0.43, 0.63], ) def process_session(drone_config: dict, output_folder: str, make_video: bool = True): sr = EXPERIMENT_PARAMS["sample_rate"] chunk_s = EXPERIMENT_PARAMS["chunk_duration_s"] chunk_n = int(sr * chunk_s) c = EXPERIMENT_PARAMS["speed_of_sound"] record_s = EXPERIMENT_PARAMS["RECORD_SECONDS"] lowcut = EXPERIMENT_PARAMS["lowcut"] highcut = EXPERIMENT_PARAMS["highcut"] az_range = EXPERIMENT_PARAMS["azimuth_range"] el_range = EXPERIMENT_PARAMS["elevation_range"] threshold = EXPERIMENT_PARAMS["angular_threshold"] use_csv = drone_config["use_csv_log"] skip_s = drone_config["skip_seconds"] corrections= drone_config["corrections_samples"] os.makedirs(output_folder, exist_ok=True) # --- GPS / CSV setup (only when drone log is available) --- ref_x = ref_y = initial_altitude = 0.0 flight_data = pd.DataFrame() transformer = None if use_csv: ref_data = pd.read_csv(drone_config["ref_csv"], low_memory=False) flight_data = pd.read_csv(drone_config["flight_csv"], low_memory=False) ref_lat = ref_data["latitude"].dropna().astype(float).mean() ref_lon = ref_data["longitude"].dropna().astype(float).mean() cols = ["latitude", "longitude", "altitude_above_seaLevel(feet)", "time(millisecond)", " xSpeed(mph)", " ySpeed(mph)", " zSpeed(mph)"] flight_data = flight_data.iloc[drone_config["start_index_csv"]:].reset_index(drop=True) flight_data = flight_data[cols].dropna() flight_data["altitude_above_seaLevel(feet)"] *= 0.3048 initial_altitude = ref_data["altitude_above_seaLevel(feet)"].iloc[0] * 0.3048 transformer = Transformer.from_crs("epsg:4326", "epsg:32756", always_xy=True) ref_x, ref_y = transformer.transform(ref_lon, ref_lat) flight_data["X_m"], flight_data["Y_m"] = transformer.transform( flight_data["longitude"].values, flight_data["latitude"].values, ) drone_init_az = calculate_azimuth_meters(ref_x, ref_y, flight_data.iloc[0]["X_m"], flight_data.iloc[0]["Y_m"]) az_offset = drone_config["initial_azimuth"] - drone_init_az el_offset = drone_config["initial_elevation"] else: az_offset = el_offset = 0.0 # --- Beamforming setup --- mic_array_cfg = drone_config.get("mic_array") or DEFAULT_MIC_ARRAY mic_array = {k: mic_array_cfg[k] for k in ("a_", "a2_", "h_", "r_") if k in mic_array_cfg} mic_positions, table_delays, _ = initialize_beamforming_params( az_range, el_range, c, sr, **mic_array ) # --- Open WAVs and seek --- wav_files = open_wav_files(drone_config["wavs_nosync"]) for i, wf in enumerate(wav_files): corr = corrections[i] if i < len(corrections) else 0 skip_wav_seconds(wf, skip_s, sr, extra_samples=corr) buffers = [np.zeros((chunk_n, EXPERIMENT_PARAMS["CHANNELS"]), dtype=np.int32) for _ in wav_files] # --- Plot setup --- plt.ion() fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(18, 8)) session_label = f"{drone_config['session']} / {drone_config['segment_name']}" fig.suptitle(session_label) ax_xy, ax_alt, ax_beam = axes[0, 0], axes[0, 1], axes[0, 2] ax_labels = axes[1, 2] for ax in axes.flat: ax.grid(True) ax_xy.set_title("XY Position (m)") ax_xy.set_xlabel("X (m)"); ax_xy.set_ylabel("Y (m)") line_xy, = ax_xy.plot([], [], "b-", label="Path") marker_xy, = ax_xy.plot([], [], "ro", label="Current") ax_xy.legend() ax_alt.set_title("Altitude vs Time") ax_alt.set_xlabel("Time (s)"); ax_alt.set_ylabel("Altitude (m)") line_alt, = ax_alt.plot([], [], "g-") marker_alt, = ax_alt.plot([], [], "mo") cax = ax_beam.imshow( np.zeros((len(el_range), len(az_range))), extent=[az_range[0], az_range[-1], el_range[0], el_range[-1]], origin="lower", aspect="auto", cmap="jet", ) fig.colorbar(cax, ax=ax_beam, label="Energy") line_csv, = ax_beam.plot([], [], "k+", markersize=30, label="GPS truth") marker_max, = ax_beam.plot([], [], "ro", label="Max energy") ax_beam.set_title("Beamforming"); ax_beam.legend() label_im = ax_labels.imshow( np.zeros((len(el_range), len(az_range))), extent=[az_range[0], az_range[-1], el_range[0], el_range[-1]], origin="lower", aspect="auto", cmap="jet", ) ax_labels.set_title("Labels") info_text = fig.text(0.84, 0.7, "", va="top", ha="left", fontsize=9, bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.5)) footer_text = fig.text(0.5, 0.02, "", va="bottom", ha="center", fontsize=10) plt.subplots_adjust(right=0.82, bottom=0.15) fig.canvas.draw(); fig.canvas.flush_events() max_iterations = min(int(sr / chunk_n * record_s), len(flight_data) if use_csv else int(sr / chunk_n * record_s)) csv_x_pos, csv_y_pos, alt_times, alt_values = [], [], [], [] results_rows = [] video_path = os.path.join(output_folder, f"{drone_config['segment_name']}.mp4") writer = FFMpegWriter(fps=10) with writer.saving(fig, video_path, dpi=100) if make_video else _null_ctx(): for time_idx in range(max_iterations): # Read audio block block = None for j, wf in enumerate(wav_files): block = read_wav_block(wf, chunk_n, EXPERIMENT_PARAMS["CHANNELS"]) if block is None: break buffers[j] = block if block is None: break combined = np.hstack(buffers) filtered = apply_bandpass_filter(combined, lowcut, highcut, sr) audio_time_s = calculate_time(time_idx, chunk_n, sr) + skip_s # Beamform energy map energy = beamform_time(filtered, table_delays) max_idx = np.unravel_index(np.argmax(energy), energy.shape) est_az = az_range[max_idx[0]] est_el = el_range[max_idx[1]] # GPS-derived ground truth (if available) if use_csv and time_idx < len(flight_data): row = flight_data.iloc[time_idx] x, y = row["X_m"], row["Y_m"] alt_abs = row["altitude_above_seaLevel(feet)"] alt_rel = alt_abs - initial_altitude vx = row[" xSpeed(mph)"] * 0.44704 vy = row[" ySpeed(mph)"] * 0.44704 vz = row[" zSpeed(mph)"] * 0.44704 csv_time_ms = row["time(millisecond)"] csv_az = wrap_angle(calculate_azimuth_meters(ref_x, ref_y, x, y) + az_offset) csv_az = -csv_az csv_el = calculate_elevation_meters(alt_abs, ref_x, ref_y, x, y, initial_altitude) + el_offset dist = calculate_total_distance_meters(ref_x, ref_y, x, y, initial_altitude, alt_abs) az_diff, el_diff = calculate_angle_difference(est_az, csv_az, est_el, csv_el) pos_x, pos_y = x - ref_x, y - ref_y else: csv_az = csv_el = csv_time_ms = dist = alt_rel = alt_abs = 0.0 vx = vy = vz = az_diff = el_diff = 0.0 pos_x = pos_y = 0.0 # Build labels grid num_az, num_el = len(az_range), len(el_range) beamformed_data = np.empty((num_az, num_el, chunk_n), dtype=np.float32) labels_grid = np.zeros((num_az, num_el), dtype=np.int32) for az_idx, grid_az in enumerate(az_range): for el_idx, grid_el in enumerate(el_range): beamformed_data[az_idx, el_idx, :] = beamform_in_direction( filtered, mic_positions, grid_az, grid_el, sr, c ) if use_csv: ang_dist = calculate_angular_distance(csv_az, csv_el, grid_az, grid_el) labels_grid[az_idx, el_idx] = 1 if ang_dist <= threshold else 0 # Save NPZ np.savez( os.path.join(output_folder, f"chunk_{time_idx:04d}.npz"), beamformed_data=beamformed_data, labels=labels_grid, csv_azimuth=csv_az, csv_elevation=csv_el, csv_time_ms=csv_time_ms, audio_time_s=audio_time_s, total_distance=dist, altitude_m=alt_abs, vx_mps=vx, vy_mps=vy, vz_mps=vz, pos_x=pos_x, pos_y=pos_y, estimated_azimuth=est_az, estimated_elevation=est_el, ) results_rows.append({ "audio_time_s": audio_time_s, "csv_time_ms": csv_time_ms, "estimated_az_deg": est_az, "estimated_el_deg": est_el, "csv_az_deg": csv_az, "csv_el_deg": csv_el, "diff_az_deg": az_diff, "diff_el_deg": el_diff, "distance_m": dist, "vx_mps": vx, "vy_mps": vy, "vz_mps": vz, }) # Update plots csv_x_pos.append(pos_x); csv_y_pos.append(pos_y) line_xy.set_data(csv_x_pos, csv_y_pos) marker_xy.set_data([pos_x], [pos_y]) ax_xy.relim(); ax_xy.autoscale_view() alt_times.append(audio_time_s); alt_values.append(alt_rel) line_alt.set_data(alt_times, alt_values) marker_alt.set_data([audio_time_s], [alt_rel]) ax_alt.relim(); ax_alt.autoscale_view() cax.set_data(energy.T) cax.set_clim(np.min(energy), np.max(energy)) marker_max.set_data([est_az], [est_el]) if use_csv: line_csv.set_data([csv_az], [csv_el]) label_im.set_data(labels_grid.T) label_im.set_clim(0, 1) info_text.set_text( f"Alt: {alt_rel:.2f} m\nDist: {dist:.2f} m\n" f"Vx:{vx:.2f} Vy:{vy:.2f} Vz:{vz:.2f} m/s" ) footer_text.set_text( f"Audio={format_time_s(audio_time_s)} | " f"SSL=(Az={est_az:.1f}, El={est_el:.1f}) | " f"CSV=(Az={csv_az:.1f}, El={csv_el:.1f}) | " f"Diff=(Az={az_diff:.1f}, El={el_diff:.1f})" ) fig.canvas.draw(); fig.canvas.flush_events() if make_video: writer.grab_frame() print( f"[{time_idx:05d}] Audio={format_time_s(audio_time_s)} | " f"SSL=(Az={est_az:.1f} El={est_el:.1f}) | " f"CSV=(Az={csv_az:.1f} El={csv_el:.1f}) | " f"Dist={dist:.1f}m" ) plt.ioff(); plt.show() for wf in wav_files: wf.close() return pd.DataFrame(results_rows) class _null_ctx: """No-op context manager used when make_video=False.""" def __enter__(self): return self def __exit__(self, *_): pass # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="Create beamformed NPZ dataset chunks from a session config." ) parser.add_argument( "--config", required=True, help="Path to alignment_params.json for the session/flight." ) parser.add_argument( "--segment", default=None, help="Segment name to process (e.g. 'drone', 'no_drone_1'). " "Required only when the config has multiple segments." ) parser.add_argument( "--output", default=None, help="Output folder for NPZ chunks. Defaults to ./npz///" ) parser.add_argument( "--no-video", action="store_true", help="Skip MP4 visualization output (faster)." ) args = parser.parse_args() cfg = load_segment_config(args.config, args.segment) output_folder = args.output or os.path.join( "npz", cfg["session"], cfg["segment_name"] ) print(f"Session : {cfg['session']}") print(f"Segment : {cfg['segment_name']}") print(f"Output : {output_folder}") print(f"CSV log : {'yes — ' + str(cfg['flight_csv']) if cfg['use_csv_log'] else 'no (no-drone mode, all labels=0)'}") print(f"skip_s : {cfg['skip_seconds']} s | corrections: {cfg['corrections_samples']}") print() df = process_session(cfg, output_folder, make_video=not args.no_video) if not df.empty: results_csv = os.path.join(output_folder, "results.csv") df.to_csv(results_csv, index=False) print(f"Results saved to {results_csv}") print("Done.") if __name__ == "__main__": main()