使用list或numpy过滤3D python列表

Question

我有一个python 3D列表。更清楚的是，一个列表的列表，其中每个列表都是一个框的四角坐标。我必须过滤所有小于某个给定大小的框。

让我们假设这是python列表。

box = [[[4, 4], [4, 8], [8, 8], [8, 4]],
      [[8, 8], [8, 16], [16, 16], [16, 8]],
      [[20,16],[20,20],[24,20],[24,16]]
      ...
      ]

我必须过滤所有长度和宽度小于或等于5的框。

filtered_box = [[[4, 4], [4, 8], [8, 8], [8, 4]],
                [[20,16],[20,20],[24,20],[24,16]]
               ...
               ]

这是我当前的代码

filtered_box = []
for c in box:
    min_x, min_y = c[0]
    max_x, max_y = c[2]
    if max_x - min_x <= 5 & min_y - max_y <= 5:
        filtered_box.append(c)

这是良好的工作，但我需要一个更优化的解决方案。它可以使用numpy并转换回python list，或者使用list上的原生python操作。我使用的是Python 3。

Answer 1

使用numpy的解决方案可能如下所示：

filtered_array = array[
    (np.abs(array[:, 0, 0] - array[:, 3, 0]) < 5) &
    (np.abs(array[:, 0, 1] - array[:, 3, 1]) < 5), :, :]

array = np.array(box)在哪里。

如果你准备了数据(numpy数组)，我猜这个解决方案会比普通的python快很多。将数据从python列表转换到numpy数组将使任何时间增益无效。

Answer 2

这个列表理解替换了你的for-loop结构：

import numpy as np

box = [
      [[4, 4], [4, 8], [8, 8], [8, 4]],
      [[8, 8], [8, 16], [16, 16], [16, 8]],
      [[20,16],[20,20],[24,20],[24,16]]
      ]

box = np.array(box) # convert to numpy array

# one-liner
filtered = [i for i in box if np.ptp(i[:,0]) <= 5 and np.ptp(i[:,1]) <= 5]

filtered = np.array(filtered) # if you want to convert to a numpy array