Python lifelines中获取Kaplan-Meier图对象句柄并设置图例

核心问题分析 #

你遇到的问题本质是：lifelines的plot方法会同时生成KM曲线和删失点的绘图对象，默认两者都会加入图例；同时循环内的图例设置时机不对，导致无法正确捕获KM曲线的handles。

修正后的代码 #

def plot_km(col, months, dpi):
    # 全局设置dpi，避免循环内重复执行
    plt.rcParams['figure.dpi'] = dpi
    plt.rcParams['savefig.dpi'] = dpi
    
    ax = plt.subplot(111)
    clrIdx = 0
    km_handles = []  # 存储两条KM曲线的绘图对象
    
    for r in df[col].unique():
        ix = df[col] == r
        kmf.fit(T[ix], C[ix], label=r)
        # 绘制KM曲线，给删失点设置特殊标签排除出图例
        ax = kmf.plot(ax=ax, color=colorsKM[clrIdx], show_censors=True,
                     censor_styles={'ms': 6, 'marker': 's', 'label': '_nolegend_'})
        # 获取刚绘制的KM曲线对象（ax.lines最后一条即为当前曲线）
        km_handles.append(ax.lines[-1])
        clrIdx += 1
    
    # 统一设置自定义图例
    ax.legend(handles=km_handles, labels=['test1', 'test2'])
    plt.show()

关键调整说明 #

1. 删失点排除图例：用matplotlib官方支持的'_nolegend_'标签替代'_out_'，确保删失点不会出现在图例中，这个标签是专门用来隐藏图例项的，兼容性更好。
2. 正确捕获曲线handles：ax.lines列表存储了当前轴上所有折线对象，每次绘制KM曲线后，最后一条就是刚生成的曲线，直接取ax.lines[-1]即可获取对应的handle。
3. 调整图例设置时机：等两条曲线都绘制完成后再调用legend，确保能拿到所有需要的handles，避免循环内提前设置导致的不生效问题。
4. 优化dpi设置：把dpi参数设置移到循环外，避免重复执行冗余操作。

原代码失效原因 #

• ax.get_legend_handles_labels()返回空：因为lifelines的plot方法默认没有正确注册曲线的图例标签，或者被删失点的标签干扰。
• 自定义图例不生效：循环内提前调用legend时，第二条曲线还未完成绘制，无法捕获完整的handles；传空handles=[]自然不会显示任何图例项。
• NameError：因为你没有定义handles变量就直接传入，导致变量未定义报错。

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