""" mmWave_RA_RE.py - 深度学习特征生成基线脚本 功能: 从雷达数据立方体生成 Range-Doppler, Range-Azimuth, Range-Elevation 特征图 输入: Cube_data 目录下的 *_cube.npy 文件 [Frames, VirtRx, Chirps, Samples] 输出: - *_features.npz: 包含 rd, ra, re 三个数组 - *_merged.npy: 拼接后的 [Frames, Range, Channels] 张量 """ import os import sys import glob import json import numpy as np from scipy.signal.windows import taylor, hamming from tqdm import tqdm # 导入配置解析器 try: from process_data import parse_radar_cfg, parse_tx_order except ImportError: try: from src.process_data import parse_radar_cfg, parse_tx_order except ImportError: sys.path.append(os.path.dirname(os.path.abspath(__file__))) from process_data import parse_radar_cfg, parse_tx_order # =================================================================== # 配置区域 # =================================================================== INPUT_CUBE_DIR = r"F:\test\Data_bin\lab\Cube_data" OUTPUT_FEATURE_DIR = r"F:\Data_bin\dormitory\Heatmap_Features_25" CFG_PATH = r"F:\test\cfg\Radar.cfg" # 算法参数 DOWNSAMPLE_FACTOR = 4 # 帧降采样因子 (1=不降采样) ANGLE_FFT_SIZE = 64 # 角度 FFT 点数 RANGE_WINDOW = "hamming" # Range FFT 窗函数 DOPPLER_WINDOW = "hamming" # Doppler FFT 窗函数 # IWR6843ISK 天线拓扑定义 IDXS_AZIMUTH = np.arange(0, 8) # 水平阵列: TX0(0-3) + TX2(4-7) # 垂直配对 (水平位置对齐): (2,8), (3,9), (4,10), (5,11) - 见 elevation_phase_diff() # =================================================================== # 窗函数缓存 (避免重复计算) # =================================================================== _WINDOW_CACHE = {} def get_window(name: str, length: int) -> np.ndarray: """获取或创建窗函数""" key = (name, length) if key not in _WINDOW_CACHE: if name == 'hamming': _WINDOW_CACHE[key] = hamming(length) elif name == 'taylor': try: _WINDOW_CACHE[key] = taylor(length, nbar=3, sll=40) except: _WINDOW_CACHE[key] = hamming(length) else: _WINDOW_CACHE[key] = np.ones(length) return _WINDOW_CACHE[key] # =================================================================== # 核心处理函数 # =================================================================== def normalize_complex_global(data, global_max=None): """保留相位的归一化""" magnitude = np.abs(data) phase = np.angle(data) denom = global_max if global_max else 1e4 return (magnitude / (denom + 1e-12)) * np.exp(1j * phase) def range_fft(cube: np.ndarray, window_type: str = 'hamming') -> np.ndarray: """ Range FFT Input: [Frames, VirtRx, Chirps, Samples] Output: [Frames, VirtRx, Chirps, RangeBins] """ n_samples = cube.shape[-1] win = get_window(window_type, n_samples) cube_win = cube * win.reshape(1, 1, 1, -1) return np.fft.fft(cube_win, axis=-1) def doppler_fft(range_cube: np.ndarray, window_type: str = 'hamming') -> np.ndarray: """ Doppler FFT Input: [Frames, VirtRx, Chirps, RangeBins] Output: [Frames, VirtRx, Doppler, RangeBins] (fftshift) """ n_chirps = range_cube.shape[2] win = get_window(window_type, n_chirps) cube_win = range_cube * win.reshape(1, 1, -1, 1) dop_fft = np.fft.fft(cube_win, axis=2) return np.fft.fftshift(dop_fft, axes=2) def clutter_removal(data: np.ndarray, axis: int = 2) -> np.ndarray: """静态杂波去除 (减去均值)""" mean_val = np.mean(data, axis=axis, keepdims=True) return data - mean_val def azimuth_fft(range_cube: np.ndarray, num_angle_bins: int = 64) -> np.ndarray: """ Azimuth FFT (生成 Range-Azimuth Map) Input: [Frames, VirtRx, Chirps, Range] Output: [Frames, Angle, Range] """ # 选取水平阵列 az_data = range_cube[:, IDXS_AZIMUTH, :, :] n_ant = az_data.shape[1] # 加窗 (Taylor 窗获得更好的旁瓣抑制) win = get_window('taylor', n_ant) az_data = az_data * win.reshape(1, -1, 1, 1) # Angle FFT angle_out = np.fft.fft(az_data, n=num_angle_bins, axis=1) angle_out = np.fft.fftshift(angle_out, axes=1) # 计算功率谱并对 Chirp 维求平均 -> [Frames, Angle, Range] ra_map = np.mean(np.abs(angle_out), axis=2) return ra_map def elevation_phase_diff(range_cube, num_angle_bins=64): """ [修正与优化] 使用4对垂直配对计算俯仰角 IWR6843ISK 物理对齐分析: - Base Array (Z=0): Index 0,1,2,3 (TX0) 和 4,5,6,7 (TX2) - Elev Array (Z=1): Index 8,9,10,11 (TX1, 物理右移2单位) 水平对齐配对 (X坐标相同): - X=2: Index 2 (TX0_RX2) & Index 8 (TX1_RX0) - X=3: Index 3 (TX0_RX3) & Index 9 (TX1_RX1) - X=4: Index 4 (TX2_RX0) & Index 10 (TX1_RX2) - X=5: Index 5 (TX2_RX1) & Index 11 (TX1_RX3) """ # [关键修正] 正确的配对列表 elevation_pairs = [(2, 8), (3, 9), (4, 10), (5, 11)] n_frames = range_cube.shape[0] # 存储每对的FFT结果 all_re_maps = [] for idx1, idx2 in elevation_pairs: # 取出这一对天线 [Frames, 2, Chirps, Range] el_data = range_cube[:, [idx1, idx2], :, :] n_ant = 2 # 每对2根天线 win = get_window('hamming', n_ant) el_data = el_data * win.reshape(1, -1, 1, 1) # 2点FFT (在 axis=1, 天线维) # 注意: 2点FFT补零到64点，利用插值获得平滑的相位响应 el_out = np.fft.fft(el_data, n=num_angle_bins, axis=1) el_out = np.fft.fftshift(el_out, axes=1) # 计算模值 (不单独归一化，避免弱信号噪声被放大) el_mag = np.abs(el_out) # 对Chirp维求平均 (非相干积累) -> [Frames, Angle, Range] re_map = np.mean(el_mag, axis=2) all_re_maps.append(re_map) # 对4对结果取平均, 提升信噪比 re_map_avg = np.mean(np.stack(all_re_maps, axis=0), axis=0) return re_map_avg # =================================================================== # 主处理流程 # =================================================================== def process_single_file(cube_path: str, output_root: str) -> None: """处理单个数据文件""" filename = os.path.basename(cube_path).replace('_cube.npy', '') # [Revert] 不再创建子目录，直接使用根目录 out_dir = output_root # os.makedirs(out_dir, exist_ok=True) # 根目录在main中已创建 # 加载数据 try: cube = np.load(cube_path) except Exception as e: tqdm.write(f" [Error] 加载 {filename} 失败: {e}") return # 数据校验 if cube.ndim != 4: tqdm.write(f" [Error] {filename} 维度错误: {cube.shape}") return n_frames, n_rx, n_chirps, n_samples = cube.shape if n_rx < 12: tqdm.write(f" [Error] {filename} 通道数不足: {n_rx} < 12") return if np.any(np.isnan(cube)) or np.any(np.isinf(cube)): tqdm.write(f" [Error] {filename} 包含 NaN 或 Inf") return # 1. Range FFT range_res = range_fft(cube, window_type=RANGE_WINDOW) # 2. Clutter Removal (MTI) range_res = clutter_removal(range_res, axis=2) # 3. Doppler FFT -> Range-Doppler Map dop_res = doppler_fft(range_res, window_type=DOPPLER_WINDOW) rd_map = np.mean(np.abs(dop_res), axis=1) # [Frames, Doppler, Range] # 4. Azimuth FFT -> Range-Azimuth Map ra_map = azimuth_fft(range_res, ANGLE_FFT_SIZE) # [Frames, Angle, Range] # 5. Elevation FFT -> Range-Elevation Map re_map = elevation_phase_diff(range_res, ANGLE_FFT_SIZE) # [Frames, Angle, Range] # 6. 降采样 if DOWNSAMPLE_FACTOR > 1: pick_indices = np.arange(0, n_frames, DOWNSAMPLE_FACTOR) rd_map = rd_map[pick_indices, :, :] ra_map = ra_map[pick_indices, :, :] re_map = re_map[pick_indices, :, :] # [Revert] 恢复旧的扁平化文件名逻辑: SampleName_features.npz / SampleName_merged.npy # 7. 保存独立特征 (NPZ) npz_path = os.path.join(out_dir, f"{filename}_features.npz") np.savez_compressed(npz_path, rd=rd_map, ra=ra_map, re=re_map) # 8. 保存合并张量 (NPY) - 适配 Transformer/3D-CNN 输入 # 转置为 [Frames, Range, Feature] rd_T = np.transpose(rd_map, (0, 2, 1)) # [F, R, D] ra_T = np.transpose(ra_map, (0, 2, 1)) # [F, R, A] re_T = np.transpose(re_map, (0, 2, 1)) # [F, R, E] # 拼接: [Frames, Range, D + A + E] merged = np.concatenate([rd_T, ra_T, re_T], axis=-1) merge_path = os.path.join(out_dir, f"{filename}_merged.npy") np.save(merge_path, merged.astype(np.float32)) # [Revert] 不再保存 metadata.json def main(): # 创建输出目录 os.makedirs(OUTPUT_FEATURE_DIR, exist_ok=True) # 获取文件列表 cube_files = sorted(glob.glob(os.path.join(INPUT_CUBE_DIR, "*_cube.npy"))) if not cube_files: print(f"未找到数据文件: {INPUT_CUBE_DIR}") return print(f"发现 {len(cube_files)} 个数据文件") print(f"输出目录: {OUTPUT_FEATURE_DIR}") print(f"降采样因子: {DOWNSAMPLE_FACTOR}, 角度FFT: {ANGLE_FFT_SIZE}") print("-" * 50) # 使用多进程并行处理 from multiprocessing import Pool, cpu_count from functools import partial num_workers = min(4, cpu_count()) # 最多4个进程 print(f"使用 {num_workers} 个并行进程") # 创建带固定参数的处理函数 process_func = partial(process_single_file, output_root=OUTPUT_FEATURE_DIR) try: with Pool(num_workers) as pool: # 使用imap显示进度 list(tqdm( pool.imap(process_func, cube_files), total=len(cube_files), desc="Processing" )) except Exception as e: print(f"并行处理出错: {e}") print("回退到串行处理...") for f in tqdm(cube_files, desc="Processing"): try: process_single_file(f, OUTPUT_FEATURE_DIR) except Exception as e: tqdm.write(f"[Error] 处理 {os.path.basename(f)} 失败: {e}") if __name__ == "__main__": main()