比较两个不同大小的列表并在 Python 中保留条件为 True 的项目？

Question

我正在编写代码以对来自对象检测模型的边界框执行非最大抑制。我有一个地理定位边界框的主列表，它们是geolocated_annotations =[[(x1, y1), (x2, y2), (x3, y3), (x4, y4)]] 坐标的嵌套列表。 (x3,y3) 和 (x4,y4) 坐标是边界框的地理位置。我在我的深度学习模型中使用了嵌入式 NMS 函数，结果是 keep=[[x1, y1, x2, y2]] 坐标的嵌套列表。 keep 中的结果没有最后两个坐标，因为嵌入函数没有考虑地理位置。

我想将 keep 列表与 geolocated_annotations 列表进行比较，并从 keep 中找到所有等于 (x1,y1) 和 (x2) 的 (x1,y1) 和 x2,y2) 坐标,y2) 在geolocated_annotations 中的坐标。来自keep 的这两个坐标与来自geolocated_annotations 的地理位置坐标相关联，因此如果 (x1,y1) 和 (x2,y2) 值相等，则 (x3,y3) 和 (x4,y4)主列表中的值可以与它们相关联。我需要得到一个包含所有四个 (x1,y1)、(x2,y2)、(x3,y3) 和 (x4,y4) 坐标的列表。

我有一个半工作代码，但最终列表中的坐标比keep 列表中的坐标更多。遍历所有坐标也需要很长时间。

两个列表的更全面的示例：

geolocated_annotations=[[(624169.99807, 4344801.75893), (624170.64013, 4344801.34655), (624159.5841700002, 4344797.0139500005), (624159.1195900001, 4344797.0139500005)], [(624159.9808900001, 4344799.4673500005), (624160.4089300002, 4344799.4673500005), (624160.4089300002, 4344798.887930001), (624159.9808900001, 4344798.887930001)]]

keep=[[624169.99807, 4344801.75893, 624170.64013, 4344801.34655],
 [624151.50361, 4344753.26513, 624151.93165, 4344752.69093],
 [624150.44917, 4344812.13107, 624150.88243, 4344811.58297]]

这是我的入门代码：

bboxes_keep=[]
for i in geolocated_annotations:
        for j in keep:
             if i[0][0]==j[0] and i[0][1]==j[1] and i[1][0]==j[2] and i[2][1]==j[3]:
                bboxes_keep.append(i)

预期输出示例：

bboxes_keep=[[(624169.99807, 4344801.75893), (624170.64013, 4344801.34655), (624159.5841700002, 4344797.0139500005), (624159.1195900001, 4344797.0139500005)]]

关于如何使我的代码更加精炼和准确的任何想法？

Answer 1

一般来说，字典查找是开始优化此类问题的最简单方法。由于您有成对的浮点数，因此在另一个 dict 中包含一个 dict 的数据存储是实现快速查找的一种方法。我的示例将此与 try/except 配对，这样，如果 x 或 y 不匹配，则会引发异常。请注意，还设置了数据存储区，因此如果您有多个具有相同 x 坐标的点，则所有可能的 y 值都将在一个键下的字典中。

geolocated_annotations=[[(624169.99807, 4344801.75893), (624170.64013, 4344801.34655), (624159.5841700002, 4344797.0139500005), (624159.1195900001, 4344797.0139500005)], [(624159.9808900001, 4344799.4673500005), (624160.4089300002, 4344799.4673500005), (624160.4089300002, 4344798.887930001), (624159.9808900001, 4344798.887930001)]]

    keep=[[624169.99807, 4344801.75893, 624170.64013, 4344801.34655],
 [624151.50361, 4344753.26513, 624151.93165, 4344752.69093],
 [624150.44917, 4344812.13107, 624150.88243, 4344811.58297]]

datastore = {}

for i in keep:

    try:
        datastore[i[0]][i[1]] = True
    except:
        datastore[i[0]] = {i[1]:True} 
    try:
        datastore[i[2]][i[3]] = True
    except:
        datastore[i[2]] = {i[3]:True}

result = []
for row in geolocated_annotations:
    for x,y in row:
        try:
            match = datastore[x][y]
            result.append([x,y])
        except:
            pass

注意：我不得不对这两个列表的结构做出一些假设，因为正如评论者所指出的，您的示例似乎有一些错误。如果我假设不正确，只需修复您的代码并让我知道。