【使用main函数，批量作图并进行图上标注】

数据处理与可视化教程

在本教程中，我们将详细介绍如何使用Python脚本批量处理CSV文件，并生成包含峰值标注的光谱图。

导入必要的库

首先，我们需要导入所需的库：

import os
import pandas as pd
import matplotlib.pyplot as plt

• os：用于处理文件和目录操作。
• pandas：用于数据处理和分析。
• matplotlib.pyplot：用于绘图。

定义函数：获取CSV文件列表

我们首先定义一个函数 get_csv_files(path)，该函数遍历指定路径下的所有文件夹，找到所有以.csv结尾的文件，并将它们的完整路径存储在列表 csv_files 中。这是为了自动化文件查找过程，避免手动指定每个文件的路径，提高代码的灵活性和可维护性：

def get_csv_files(path):
    csv_files = []
    try:
        for root, dirs, files in os.walk(path):
            for file in files:
                if file.lower().endswith('.csv'):
                    csv_files.append(os.path.join(root, file))
    except UnicodeDecodeError as e:
        print(f"Error reading directory {path}: {e}")
    return csv_files

定义函数：处理单个文件并生成图表

接下来，我们定义 process_file(file) 函数，该函数负责读取CSV文件、生成光谱图并标注峰值：

1. 读取CSV文件：使用 pandas 库的 read_csv 函数读取CSV文件内容，以便进行数据处理。
2. 绘图：使用 matplotlib 库的 plot 函数绘制光谱图，并设置坐标轴标签。
3. 标注峰值：通过 y.idxmax() 找到光谱强度的峰值位置，并使用 annotate 函数在图中标注峰值。
4. 保存图像：根据文件名和目录结构生成图片名称，并保存绘制的图像。这可以帮助我们直观地查看每个CSV文件中的光谱数据，并标注出重要的峰值信息。

def process_file(file):
    try:
        df = pd.read_csv(file, encoding='latin1')
        plt.figure(figsize=(15, 8))
        x = df.iloc[:, 0]
        y = df.iloc[:, 1]
        plt.plot(x, y)
        plt.xlabel('波长')
        plt.ylabel('光谱强度（dB）')
        
        # 找到峰值位置并标注
        peak_idx = y.idxmax()
        peak_x = x.iloc[peak_idx]
        peak_y = y.iloc[peak_idx]
        plt.annotate(f'峰值: {peak_y:.2f} dB', xy=(peak_x, peak_y), xytext=(peak_x + 10, peak_y - 1),
                     arrowprops=dict(facecolor='black', shrink=0.1, width=0.5, headwidth=8))
        
        base_name = os.path.basename(file)
        dir_name = os.path.dirname(file)
        sub_dir_name = os.path.basename(dir_name)
        
        # 组合目录名和文件名的一部分作为图片名称
        plot_name = '选取' + os.path.join(base_name, sub_dir_name).replace('.csv', '').replace('\', '编号')
        plt.title(plot_name, pad=20)
        plt.savefig(f'{plot_name}.png', bbox_inches='tight')
        plt.close()
    except Exception as e:
        print(f"Error processing {file}: {e}")

定义主函数：处理所有文件

然后，我们定义 main(directory) 函数，该函数负责调用上述函数来处理指定目录中的所有CSV文件：

def main(directory):
    csv_files = get_csv_files(directory)
    for file in csv_files:
        process_file(file)

为什么需要这个函数？

使用 main 函数可以让代码结构更加清晰，并且更容易扩展和维护。在需要处理不同目录或添加其他处理逻辑时，只需修改 main 函数中的相应部分即可。

1. 获取CSV文件列表：通过调用 get_csv_files 函数获取指定目录中的所有CSV文件。
2. 处理每个文件：通过调用 process_file 函数依次处理每个CSV文件，生成相应的光谱图并保存。

运行主函数

最后，我们在脚本的主程序中调用 main 函数，指定要处理的目录路径：

if __name__ == '__main__':
    # 示例调用
    path = r'测试数据'
    main(path)

为什么使用 if __name__ == '__main__'？

这一行代码是为了确保当脚本作为模块被导入时，不会自动执行 main 函数。只有在直接运行脚本时，才会调用 main 函数。这种做法有助于提高代码的复用性和模块化。

总结

通过本教程，我们学习了如何编写Python脚本来批量处理CSV文件，并生成包含峰值标注的光谱图。希望这些步骤能够帮助您更好地理解和应用这些技术。