如何构建Python交互GUI：点击调整绘图标记位置

需求背景 #

使用pcolormesh绘制回声数据集，需要用红色'x'标记图中的亮斑。但由于数据与标签采集方式问题，原有标签坐标常与实际亮斑位置不符。需要构建工具实现：

• 生成带有原始标签的初始绘图
• 支持用户点击修正标记位置（点击新位置替换旧标记）
• 最终保存修正后的图像与新标签坐标

修改后的完整代码 #

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from matplotlib.backend_bases import MouseButton

# 假设ECHOMAP_PIXELS和N是预先定义的全局变量，可根据实际调整
ECHOMAP_PIXELS = 800
N = 100

def echomap_plot_interactive(d0df, d1df, t0, t1, scan, save_to, keys, labels=None):
    """
    交互式回声图绘制工具，支持点击修正标记位置并保存结果
    
    Inputs:
        d0df: pandas dataframe of data
        d1df: pandas dataframe of data
        t0: x values for d0df
        t1: x values for d1df
        scan: scan number
        save_to: path to save echomap plots and labels to
        keys: list of frequencies (220000, 291000)
        labels: list of tuples of coordinates to label on the plot.
    """
    my_dpi = 96
    
    d0 = d0df.to_numpy()
    d1 = d1df.to_numpy()
    r0 = np.linspace(0,50,N)
    r1 = np.linspace(0,50,N)

    # 初始化标签列表，为空则创建空列表
    if labels is None:
        labels = []
    else:
        labels = list(labels)
    
    # 创建绘图
    fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, 
                                   figsize=(2*(ECHOMAP_PIXELS/my_dpi), ECHOMAP_PIXELS/my_dpi), 
                                   dpi=my_dpi)
    vm = int(np.maximum(np.max(d0), np.max(d1))) + 1

    im = ax0.pcolormesh(t0, r0, d0, vmin=0, vmax=vm)
    fig.colorbar(im, ax=ax0)
    ax0.set_title(f'Echomap Plot of 50m (Long Range) Scan {scan} {keys[0]} Data')
    ax0.grid()

    im1 = ax1.pcolormesh(t1, r1, d1, vmin=0, vmax=vm)
    fig.colorbar(im1, ax=ax1)
    ax1.set_title(f'Echomap Plot of 50m (Long Range) Scan {scan} {keys[1]} Data')
    ax1.grid()
    
    # 存储标记点的艺术家对象，方便后续更新
    marker_artists = []
    def update_markers():
        # 清除旧标记
        for artist in marker_artists:
            artist.remove()
        marker_artists.clear()
        # 绘制新标记
        if labels:
            x_values, y_values = zip(*labels)
            # 左图标记
            artists0 = ax0.scatter(x_values, y_values, marker='x', color='red', s=250, linewidth=2)
            marker_artists.append(artists0)
            # 右图标记
            artists1 = ax1.scatter(x_values, y_values, marker='x', color='red', s=250, linewidth=2)
            marker_artists.append(artists1)
        fig.canvas.draw_idle()
    
    # 初始绘制标记
    update_markers()

    fig.suptitle(f'Scan {scan} Echomaps (点击修正标记，按S保存退出)')
    plt.tight_layout()

    # 点击事件处理函数
    def on_click(event):
        # 只处理左键点击，且点击在坐标轴内
        if event.button is not MouseButton.LEFT or not event.inaxes:
            return
        
        # 获取点击的坐标
        click_x = event.xdata
        click_y = event.ydata
        
        # 找到距离点击位置最近的标记并替换；无标记则添加新标记
        if labels:
            distances = [np.hypot(x - click_x, y - click_y) for x, y in labels]
            min_dist_idx = np.argmin(distances)
            labels[min_dist_idx] = (click_x, click_y)
        else:
            labels.append((click_x, click_y))
        
        # 更新绘图标记
        update_markers()
    
    # 键盘事件处理函数（保存并退出）
    def on_key(event):
        if event.key.lower() == 's':
            # 创建保存目录
            if not os.path.exists(save_to):
                os.makedirs(save_to)
            
            # 保存图像
            img_path = os.path.join(save_to, f"{scan}_corrected_echomaps.png")
            fig.savefig(img_path)
            print(f"修正后的图像已保存至: {img_path}")
            
            # 保存标签坐标到CSV文件
            labels_df = pd.DataFrame(labels, columns=['x', 'y'])
            label_path = os.path.join(save_to, f"{scan}_corrected_labels.csv")
            labels_df.to_csv(label_path, index=False)
            print(f"修正后的标签已保存至: {label_path}")
            
            # 关闭绘图窗口
            plt.close(fig)
    
    # 绑定事件
    fig.canvas.mpl_connect('button_press_event', on_click)
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    # 显示绘图
    plt.show()

关键功能说明 #

• 标记更新机制：通过update_markers函数统一管理标记的绘制与清除，确保每次点击后标记实时更新
• 交互逻辑：
  • 左键点击绘图区域：若已有标记，替换距离点击位置最近的标记；若无标记则添加新标记
  • 按S键：保存修正后的图像与标签坐标（标签保存为CSV格式）
• 兼容性：保留了原函数的输入参数结构，可直接替换原函数使用

使用示例 #

# 假设已有d0df, d1df, t0, t1等数据
# 原始标签示例
original_labels = [(10, 20), (30, 40)]

# 调用交互式工具
echomap_plot_interactive(
    d0df=d0df,
    d1df=d1df,
    t0=t0,
    t1=t1,
    scan=1,
    save_to="./corrected_echomaps/",
    keys=[220000, 291000],
    labels=original_labels
)

内容的提问来源于stack exchange，提问作者sophivar