import os import cv2 import matplotlib import numpy as np import argparse import pandas as pd from bil2rgb import process_bil_files from mask import detect_microplastic_mask_from_array from shape_spectral import process_images from shape_spectral_background import process_images_background import time from tqdm import tqdm matplotlib.use('TkAgg') # 训练相机波长（237通道） TRAIN_WAVELENGTHS = [912.36, 915.68, 919, 922.31, 925.63, 928.95, 932.27, 935.59, 938.91, 942.23, 945.55, 948.87, 952.18, 955.5, 958.82, 962.14, 965.46, 968.78, 972.1, 975.42, 978.74, 982.06, 985.38, 988.7, 992.02, 995.34, 998.65, 1002, 1005.3, 1008.6, 1011.9, 1015.3, 1018.6, 1021.9, 1025.2, 1028.5, 1031.9, 1035.2, 1038.5, 1041.8, 1045.1, 1048.5, 1051.8, 1055.1, 1058.4, 1061.7, 1065.1, 1068.4, 1071.7, 1075, 1078.3, 1081.7, 1085, 1088.3, 1091.6, 1094.9, 1098.3, 1101.6, 1104.9, 1108.2, 1111.5, 1114.9, 1118.2, 1121.5, 1124.8, 1128.1, 1131.5, 1134.8, 1138.1, 1141.4, 1144.8, 1148.1, 1151.4, 1154.7, 1158, 1161.4, 1164.7, 1168, 1171.3, 1174.6, 1178, 1181.3, 1184.6, 1187.9, 1191.3, 1194.6, 1197.9, 1201.2, 1204.5, 1207.9, 1211.2, 1214.5, 1217.8, 1221.2, 1224.5, 1227.8, 1231.1, 1234.4, 1237.8, 1241.1, 1244.4, 1247.7, 1251.1, 1254.4, 1257.7, 1261, 1264.3, 1267.7, 1271, 1274.3, 1277.6, 1281, 1284.3, 1287.6, 1290.9, 1294.2, 1297.6, 1300.9, 1304.2, 1307.5, 1310.9, 1314.2, 1317.5, 1320.8, 1324.2, 1327.5, 1330.8, 1334.1, 1337.4, 1340.8, 1344.1, 1347.4, 1350.7, 1354.1, 1357.4, 1360.7, 1364, 1367.4, 1370.7, 1374, 1377.3, 1380.7, 1384, 1387.3, 1390.6, 1394, 1397.3, 1400.6, 1403.9, 1407.2, 1410.6, 1413.9, 1417.2, 1420.5, 1423.9, 1427.2, 1430.5, 1433.8, 1437.2, 1440.5, 1443.8, 1447.1, 1450.5, 1453.8, 1457.1, 1460.4, 1463.8, 1467.1, 1470.4, 1473.7, 1477.1, 1480.4, 1483.7, 1487, 1490.4, 1493.7, 1497, 1500.3, 1503.7, 1507, 1510.3, 1513.6, 1517, 1520.3, 1523.6, 1526.9, 1530.3, 1533.6, 1536.9, 1540.2, 1543.6, 1546.9, 1550.2, 1553.6, 1556.9, 1560.2, 1563.5, 1566.9, 1570.2, 1573.5, 1576.8, 1580.2, 1583.5, 1586.8, 1590.1, 1593.5, 1596.8, 1600.1, 1603.4, 1606.8, 1610.1, 1613.4, 1616.7, 1620.1, 1623.4, 1626.7, 1630.1, 1633.4, 1636.7, 1640, 1643.4, 1646.7, 1650, 1653.3, 1656.7, 1660, 1663.3, 1666.7, 1670, 1673.3, 1676.6, 1680, 1683.3, 1686.6, 1689.9, 1693.3, 1696.6, 1699.9, 1703.3, 1706.6] # 微塑料类型映射 PLASTIC_TYPE_MAPPING = { 'ABS': 0, 'HDPE': 1, 'LDPE': 2, 'PA6': 3, 'PET': 4, 'PP': 5, 'PS': 6, 'PTFE': 7, 'PVC': 8, } def get_plastic_label_from_filename(filename): """ 从文件名中提取微塑料类型并返回对应的数字标签 只要文件名中包含微塑料名称就返回对应的数字 """ filename_upper = filename.upper() for plastic_type, label in PLASTIC_TYPE_MAPPING.items(): if plastic_type in filename_upper: return label return None def read_wavelengths_from_hdr(bil_path): hdr_path = os.path.splitext(bil_path)[0] + '.hdr' if not os.path.exists(hdr_path): return np.array([], dtype=np.float64) with open(hdr_path, 'r') as f: txt = f.read() if 'wavelength' not in txt: return np.array([], dtype=np.float64) seg = txt.split('wavelength', 1)[1] seg = seg[seg.find('{')+1: seg.find('}')] vals = [v.strip() for v in seg.split(',') if v.strip()] try: waves = np.array([float(v) for v in vals], dtype=np.float64) except Exception: waves = np.array([], dtype=np.float64) return waves def resample_spectra_matrix(X, src_waves, dst_waves): src = np.asarray(src_waves, dtype=np.float64) dst = np.asarray(dst_waves, dtype=np.float64) X = np.asarray(X, dtype=np.float64) if src.size == 0 or dst.size == 0: return X # 线性插值，越界用端点外推，避免维度缺失 out = np.empty((X.shape[0], dst.size), dtype=np.float64) for i in range(X.shape[0]): row = X[i] out[i] = np.interp(dst, src, row, left=row[0], right=row[-1]) return out def read_hdr_file(bil_path): hdr_path = bil_path.replace('.bil', '.hdr') with open(hdr_path, 'r') as f: header = f.readlines() samples, lines = None, None for line in header: if line.startswith('samples'): samples = int(line.split('=')[-1].strip()) if line.startswith('lines'): lines = int(line.split('=')[-1].strip()) return samples, lines def change_hdr_file(bil_path, wavelengths=None): hdr_path = os.path.splitext(bil_path)[0] + '.hdr' if not os.path.exists(hdr_path): print(f"错误: 找不到对应的HDR文件: {hdr_path}") return # 仅在缺少 wavelength 字段时才尝试写入 with open(hdr_path, 'r', encoding='utf-8', errors='ignore') as file: content = file.read() if 'wavelength' in content: print(f"{os.path.basename(hdr_path)} 已包含 wavelength 字段，跳过追加。") return if wavelengths is None or len(wavelengths) == 0: print("HDR 缺少 wavelength，但未提供 wavelengths，跳过写入以避免错误。") return needs_newline = not content.endswith('\n') wavelength_info = "wavelength = {" + ", ".join(str(float(w)) for w in wavelengths) + "}\n" with open(hdr_path, 'a', encoding='utf-8', errors='ignore') as file: if needs_newline: file.write('\n') file.write(wavelength_info) print(f"已在 {os.path.basename(hdr_path)} 末尾追加 wavelength 字段。") def validate_inputs(bil_path, output_dir): """验证输入文件和参数""" # 检查BIL和HDR文件存在 if not os.path.exists(bil_path): raise FileNotFoundError(f"BIL文件不存在: {bil_path}") hdr_path = os.path.splitext(bil_path)[0] + '.hdr' if not os.path.exists(hdr_path): raise FileNotFoundError(f"HDR文件不存在: {hdr_path}") # 检查输出目录可写 if output_dir and not os.path.exists(output_dir): try: os.makedirs(output_dir, exist_ok=True) except Exception as e: raise RuntimeError(f"无法创建输出目录: {output_dir}") from e # 检查BIL文件波段数是否足够 try: from spectral.io import envi img = envi.open(hdr_path, bil_path) n_bands = img.nbands # bil2rgb需要波段索引9, 59, 159 if n_bands <= 159: raise ValueError(f"BIL文件波段数不足: 需要至少160个波段，但只有{n_bands}个") except Exception as e: raise RuntimeError(f"无法读取BIL文件头信息: {bil_path}") from e def generate_rgb(bil_path): """处理BIL文件生成RGB图像""" try: rgb_img = process_bil_files(bil_path) return rgb_img except Exception as e: raise RuntimeError(f"生成RGB图像失败: bil_path={bil_path}") from e def run_segmentation(rgb_img, segmentation_model_path=None): """运行分割获取掩膜""" try: mask, filter_mask_original = detect_microplastic_mask_from_array( image=rgb_img, filter_method='threshold', diameter=None, flow_threshold=0.4, cellprob_threshold=-1, model_path=segmentation_model_path, detect_filter=True ) return mask, filter_mask_original except Exception as e: raise RuntimeError("分割失败: 无法检测微塑料颗粒") from e def extract_primary_features(bil_path, mask): """提取主要特征""" try: df = process_images(bil_path, mask) return df except Exception as e: raise RuntimeError(f"特征提取失败: bil_path={bil_path}") from e def compute_background_spectrum(bil_path, mask): """计算背景光谱""" try: bg_spectrum = process_images_background(bil_path, mask) return bg_spectrum except Exception as e: raise RuntimeError(f"背景光谱计算失败: bil_path={bil_path}") from e def apply_background_correction(df, bg_spectrum): """应用背景校正，不进行重采样""" # 识别光谱列：所有以wavelength_开头的列 spec_cols = [c for c in df.columns if isinstance(c, str) and c.startswith('wavelength_')] if not spec_cols: raise ValueError("未找到光谱列（以wavelength_开头的列）") # 创建原始光谱数据的副本 df_original = df.copy() # 背景校正：用背景光谱逐列相除 bg = np.asarray(bg_spectrum, dtype=np.float64).ravel() # 尾部对齐，取最小长度，避免维度不一致 n = min(len(spec_cols), bg.shape[0]) use_cols = spec_cols[-n:] df_corrected = df.copy() df_corrected.loc[:, use_cols] = df_corrected.loc[:, use_cols].div(bg[-n:], axis=1) return df_corrected, df_original, bg def clean_and_select_columns(df): """数据清理和列选择""" # 移除NaN值 df = df.dropna() # 过滤轮廓点数不足的样本 df = df[df['contour'].apply(lambda x: len(x) > 1 if isinstance(x, list) else True)] # 过滤面积过小的样本 df = df[df['area'] >= 500] return df def rename_wavelength_columns(df, prefix=''): """ 将列名中的 'wavelength_' 前缀移除，替换为指定前缀或直接波长数值 """ new_columns = {} for col in df.columns: if isinstance(col, str) and col.startswith('wavelength_'): # 提取波长数值 wavelength_value = col.replace('wavelength_', '') new_columns[col] = wavelength_value if new_columns: df = df.rename(columns=new_columns) return df def save_spectra_to_csv(df_corrected, df_original, bg_spectrum, bil_path, output_dir, plastic_label, all_corrected_data, all_original_data, all_background_data): """ 保存三种光谱数据为CSV文件 - 背景校正后的光谱 - 原始光谱 - 背景光谱 同时收集数据用于统一合并 """ base_name = os.path.splitext(os.path.basename(bil_path))[0] # 创建输出子目录 corrected_dir = os.path.join(output_dir, 'corrected_spectra') original_dir = os.path.join(output_dir, 'original_spectra') background_dir = os.path.join(output_dir, 'background_spectra') os.makedirs(corrected_dir, exist_ok=True) os.makedirs(original_dir, exist_ok=True) os.makedirs(background_dir, exist_ok=True) # 获取光谱列 spec_cols = [c for c in df_corrected.columns if isinstance(c, str) and c.startswith('wavelength_')] # 移除波长列名中的 'wavelength_' 前缀 df_corrected_renamed = rename_wavelength_columns(df_corrected.copy()) df_original_renamed = rename_wavelength_columns(df_original.copy()) # 获取新的波长列名（已移除前缀） wavelength_cols = [c for c in df_corrected_renamed.columns if c not in [col for col in df_corrected.columns if isinstance(col, str) and not col.startswith('wavelength_')]] # 保存背景校正后的光谱 df_corrected_out = df_corrected_renamed.copy() if plastic_label is not None: if len(df_corrected_out) > 0: non_spec_cols = [c for c in df_corrected_out.columns if c not in wavelength_cols] if non_spec_cols: first_col = non_spec_cols[0] df_corrected_out[first_col] = plastic_label # 添加文件名列用于区分来源 df_corrected_out.insert(0, 'source_file', base_name) corrected_path = os.path.join(corrected_dir, f"{base_name}_corrected.csv") df_corrected_out.to_csv(corrected_path, index=False) print(f" 背景校正光谱已保存: {corrected_path}") # 收集到合并列表 all_corrected_data.append(df_corrected_out) # 保存原始光谱 df_original_out = df_original_renamed.copy() if plastic_label is not None: if len(df_original_out) > 0: non_spec_cols = [c for c in df_original_out.columns if c not in wavelength_cols] if non_spec_cols: first_col = non_spec_cols[0] df_original_out[first_col] = plastic_label # 添加文件名列用于区分来源 df_original_out.insert(0, 'source_file', base_name) original_path = os.path.join(original_dir, f"{base_name}_original.csv") df_original_out.to_csv(original_path, index=False) print(f" 原始光谱已保存: {original_path}") # 收集到合并列表 all_original_data.append(df_original_out) # 保存背景光谱 # 移除 'wavelength_' 前缀 wavelength_names = [col.replace('wavelength_', '') for col in spec_cols[-len(bg_spectrum):]] if len(spec_cols) >= len(bg_spectrum) else [col.replace('wavelength_', '') for col in spec_cols] bg_df = pd.DataFrame({ 'wavelength': wavelength_names, 'background_value': bg_spectrum }) bg_df.insert(0, 'source_file', base_name) bg_df.insert(1, 'plastic_type', plastic_label if plastic_label is not None else 'unknown') background_path = os.path.join(background_dir, f"{base_name}_background.csv") bg_df.to_csv(background_path, index=False) print(f" 背景光谱已保存: {background_path}") # 收集到合并列表 all_background_data.append(bg_df) def save_combined_csv(all_corrected_data, all_original_data, all_background_data, output_dir): """ 将所有收集的数据合并保存为统一的CSV文件 """ combined_dir = os.path.join(output_dir, 'combined') os.makedirs(combined_dir, exist_ok=True) # 合并背景校正光谱 if all_corrected_data: combined_corrected = pd.concat(all_corrected_data, ignore_index=True) corrected_combined_path = os.path.join(combined_dir, 'all_corrected_spectra.csv') combined_corrected.to_csv(corrected_combined_path, index=False) print(f"\n 合并背景校正光谱已保存: {corrected_combined_path}") print(f" 总行数: {len(combined_corrected)}") # 合并原始光谱 if all_original_data: combined_original = pd.concat(all_original_data, ignore_index=True) original_combined_path = os.path.join(combined_dir, 'all_original_spectra.csv') combined_original.to_csv(original_combined_path, index=False) print(f" 合并原始光谱已保存: {original_combined_path}") print(f" 总行数: {len(combined_original)}") # 合并背景光谱 if all_background_data: combined_background = pd.concat(all_background_data, ignore_index=True) background_combined_path = os.path.join(combined_dir, 'all_background_spectra.csv') combined_background.to_csv(background_combined_path, index=False) print(f" 合并背景光谱已保存: {background_combined_path}") print(f" 总行数: {len(combined_background)}") def process_single_bil(bil_path, output_dir, segmentation_model_path=None, all_corrected_data=None, all_original_data=None, all_background_data=None): """处理单个BIL文件""" try: print(f"\n处理文件: {bil_path}") # 从文件名获取微塑料标签 filename = os.path.basename(bil_path) plastic_label = get_plastic_label_from_filename(filename) if plastic_label is not None: print(f" 检测到微塑料类型: {list(PLASTIC_TYPE_MAPPING.keys())[list(PLASTIC_TYPE_MAPPING.values()).index(plastic_label)]} -> {plastic_label}") else: print(f" 警告: 无法从文件名识别微塑料类型") # 验证输入 validate_inputs(bil_path, output_dir) bands = [912.36, 915.68, 919, 922.31, 925.63, 928.95, 932.27, 935.59, 938.91, 942.23, 945.55, 948.87, 952.18, 955.5, 958.82, 962.14, 965.46, 968.78, 972.1, 975.42, 978.74, 982.06, 985.38, 988.7, 992.02, 995.34, 998.65, 1002, 1005.3, 1008.6, 1011.9, 1015.3, 1018.6, 1021.9, 1025.2, 1028.5, 1031.9, 1035.2, 1038.5, 1041.8, 1045.1, 1048.5, 1051.8, 1055.1, 1058.4, 1061.7, 1065.1, 1068.4, 1071.7, 1075, 1078.3, 1081.7, 1085, 1088.3, 1091.6, 1094.9, 1098.3, 1101.6, 1104.9, 1108.2, 1111.5, 1114.9, 1118.2, 1121.5, 1124.8, 1128.1, 1131.5, 1134.8, 1138.1, 1141.4, 1144.8, 1148.1, 1151.4, 1154.7, 1158, 1161.4, 1164.7, 1168, 1171.3, 1174.6, 1178, 1181.3, 1184.6, 1187.9, 1191.3, 1194.6, 1197.9, 1201.2, 1204.5, 1207.9, 1211.2, 1214.5, 1217.8, 1221.2, 1224.5, 1227.8, 1231.1, 1234.4, 1237.8, 1241.1, 1244.4, 1247.7, 1251.1, 1254.4, 1257.7, 1261, 1264.3, 1267.7, 1271, 1274.3, 1277.6, 1281, 1284.3, 1287.6, 1290.9, 1294.2, 1297.6, 1300.9, 1304.2, 1307.5, 1310.9, 1314.2, 1317.5, 1320.8, 1324.2, 1327.5, 1330.8, 1334.1, 1337.4, 1340.8, 1344.1, 1347.4, 1350.7, 1354.1, 1357.4, 1360.7, 1364, 1367.4, 1370.7, 1374, 1377.3, 1380.7, 1384, 1387.3, 1390.6, 1394, 1397.3, 1400.6, 1403.9, 1407.2, 1410.6, 1413.9, 1417.2, 1420.5, 1423.9, 1427.2, 1430.5, 1433.8, 1437.2, 1440.5, 1443.8, 1447.1, 1450.5, 1453.8, 1457.1, 1460.4, 1463.8, 1467.1, 1470.4, 1473.7, 1477.1, 1480.4, 1483.7, 1487, 1490.4, 1493.7, 1497, 1500.3, 1503.7, 1507, 1510.3, 1513.6, 1517, 1520.3, 1523.6, 1526.9, 1530.3, 1533.6, 1536.9, 1540.2, 1543.6, 1546.9, 1550.2, 1553.6, 1556.9, 1560.2, 1563.5, 1566.9, 1570.2, 1573.5, 1576.8, 1580.2, 1583.5, 1586.8, 1590.1, 1593.5, 1596.8, 1600.1, 1603.4, 1606.8, 1610.1, 1613.4, 1616.7, 1620.1, 1623.4, 1626.7, 1630.1, 1633.4, 1636.7, 1640, 1643.4, 1646.7, 1650, 1653.3, 1656.7, 1660, 1663.3, 1666.7, 1670, 1673.3, 1676.6, 1680, 1683.3, 1686.6, 1689.9, 1693.3, 1696.6, 1699.9, 1703.3, 1706.6] # 修改HDR文件 change_hdr_file(bil_path, bands) # 处理BIL文件生成RGB图像 print(" 生成RGB图像...") rgb_img = generate_rgb(bil_path) # 分割阶段 print(" 生成掩膜...") mask, filter_mask_original = run_segmentation(rgb_img, segmentation_model_path) # 提取特征 print(" 提取光谱特征...") df = extract_primary_features(bil_path, mask) # 背景校正 print(" 计算背景光谱并应用校正...") bg_spectrum = compute_background_spectrum(bil_path, mask) # 应用背景校正（不进行重采样） df_corrected, df_original, bg_spectrum_aligned = apply_background_correction(df, bg_spectrum) # 数据清理 print(" 清理数据...") df_corrected = clean_and_select_columns(df_corrected) df_original = clean_and_select_columns(df_original) # 保存三种光谱数据（同时收集到合并列表） save_spectra_to_csv(df_corrected, df_original, bg_spectrum_aligned, bil_path, output_dir, plastic_label, all_corrected_data, all_original_data, all_background_data) print(f" 处理完成: {filename}") return True except Exception as e: print(f" 处理失败: {bil_path} - {e}") return False def parse_arguments(): """解析命令行参数""" parser = argparse.ArgumentParser( description='批量处理高光谱图像，提取并保存背景校正光谱、原始光谱和背景光谱' ) # 输入文件夹（包含BIL文件） parser.add_argument('--input_dir', required=True, help='包含输入BIL文件的文件夹路径') # 输出文件夹 parser.add_argument('--output_dir', required=True, help='保存CSV输出结果的文件夹路径') return parser.parse_args() def main(): """主函数""" args = parse_arguments() input_dir = args.input_dir output_dir = args.output_dir # 记录总开始时间 total_start_time = time.time() # 检查输入目录 if not os.path.exists(input_dir): print(f"错误: 输入目录不存在: {input_dir}") return # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 获取所有BIL文件 bil_files = [f for f in os.listdir(input_dir) if f.endswith('.bil')] bil_files.sort() if not bil_files: print(f"警告: 在 {input_dir} 中未找到BIL文件") return print(f"\n{'=' * 60}") print(f"找到 {len(bil_files)} 个BIL文件需要处理") print(f"{'=' * 60}") # 用于收集所有数据的列表 all_corrected_data = [] all_original_data = [] all_background_data = [] # 使用tqdm显示进度条 success_count = 0 fail_count = 0 for bil_file in tqdm(bil_files, desc="处理进度", unit="文件"): bil_path = os.path.join(input_dir, bil_file) if process_single_bil(bil_path, output_dir, all_corrected_data=all_corrected_data, all_original_data=all_original_data, all_background_data=all_background_data): success_count += 1 else: fail_count += 1 # 保存合并的CSV文件 if all_corrected_data or all_original_data or all_background_data: print(f"\n{'=' * 60}") print("正在生成合并的CSV文件...") save_combined_csv(all_corrected_data, all_original_data, all_background_data, output_dir) # 计算总耗时 total_time = time.time() - total_start_time # 打印总结 print(f"\n{'=' * 60}") print("处理完成总结") print(f"{'=' * 60}") print(f"成功处理: {success_count} 个文件") print(f"失败: {fail_count} 个文件") print(f"总耗时: {total_time:.2f} 秒") print(f"平均每个文件: {total_time / len(bil_files):.2f} 秒") print(f"{'=' * 60}") print(f"结果已保存至: {output_dir}") print(f" - 单独文件:") print(f" - 背景校正光谱: {os.path.join(output_dir, 'corrected_spectra')}") print(f" - 原始光谱: {os.path.join(output_dir, 'original_spectra')}") print(f" - 背景光谱: {os.path.join(output_dir, 'background_spectra')}") print(f" - 合并文件:") print(f" - 合并后的光谱数据: {os.path.join(output_dir, 'combined')}") if __name__ == "__main__": main()