优化Pandas DataFrame迭代查询：高效构建标注字典并关联元数据

现有如下格式的CSV文件：

Detection,Imagename,Frame_Identifier,TL_x,TL_y,BR_x,BR_y,detection_Confidence,Target_Length,Species,Confidence
0,201503.20150619.181140817.204628.jpg,0,272,142.375,382.5,340,0.475837,0,fish,0.475837
1,201503.20150619.181141498.204632.jpg,3,267.75,6.375,422.875,80.75,0.189145,0,fish,0.189145
2,201503.20150619.181141662.204633.jpg,4,820.25,78.625,973.25,382.5,0.615788,0,fish,0.615788
3,201503.20150619.181141662.204633.jpg,4,1257,75,1280,116,0.307278,0,fish,0.307278
4,201503.20150619.181141834.204634.jpg,5,194,281,233,336,0.586944,0,fish,0.586944

已将其加载为pandas.DataFrame（命名为imageannotation），需要生成以Imagename为键的字典，值为包含['bbox', 'species']的子字典：

• bbox：由TL_x, TL_y, BR_x, BR_y组成的数组
• species：对应Species列的数组

当前实现代码如下，但处理大文件时速度极慢：

test = {
    i: {
        "bbox": imageannotation[imageannotation["Imagename"] == i][
            ["TL_x", "TL_y", "BR_x", "BR_y"]
        ].values,
        "species": imageannotation[imageannotation["Imagename"] == i][
            ["Species"]
        ].values,
    }
    for i in imageannotation["Imagename"].unique()
}

最终目标是给包含唯一Imagename的imagemetadata DataFrame添加annotation列，当前用循环实现：

for i in mydict:
    imagemetadata.loc[imagemetadata.Imagename == i, "annotation"] = [test[i]]

1. 优化字典/结构化数据生成 #

利用pandas的groupby分组操作替代循环切片，groupby是内部优化的向量化操作，避免多次重复查询DataFrame的开销：

# 按Imagename分组，生成每个图片对应的注释结构
annotations_series = imageannotation.groupby('Imagename').apply(
    lambda group: {
        'bbox': group[['TL_x', 'TL_y', 'BR_x', 'BR_y']].values,
        'species': group[['Species']].values
    }
)

如果需要保留字典格式，直接转换即可：

test_dict = annotations_series.to_dict()

2. 快速为imagemetadata添加列 #

放弃循环逐行赋值，改用map向量式操作匹配Imagename，效率远高于循环loc：

# 直接通过Imagename映射注释数据到新列
imagemetadata['annotation'] = imagemetadata['Imagename'].map(annotations_series)

优势说明 #

• groupby仅需对imageannotation做一次分组遍历，避免了原方法中对每个唯一Imagename重复切片的冗余操作，大数据量下速度提升显著
• map是pandas的向量式操作，内部基于C实现，比Python循环逐行赋值效率高几个数量级

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