python - 在 Numpy 中查找 Top10(n) RGB 颜色

Question

我有一个吃几秒钟的功能。该函数应返回给定图像中的 Top(n) 颜色。必须对返回进行排序，因此我可以使用第一个、第二个、第三个最顶部颜色的 rgb 值。

第一手我有一个 PIL.Image 对象，我在 x,y 坐标上循环并将其计算在 defaultdict 中。我已经用 Numpy 数组替换了项目中的 PIL 对象，这给了我很大的提升，但我不知道在这种情况下如何替换 defaultdict。

我目前的解决方案：

import numpy as np
from scipy import misc  # for example Image
from collections import defaultdict

def count_colors(img, n):
    img = img.reshape(-1, img.shape[-1])
    color = defaultdict(int)

    for pixel in img:
        rgb = (pixel[0], pixel[1], pixel[2])
        color[rgb] += 1

    sorted_color = sorted(color.items(), key=lambda k_v: k_v[1], reverse=True)
    sorted_color = sorted_color[:n]

    return sorted_color

img = misc.face()  # example Numpy Image array
top_colors = count_colors(img, n=5)

display(top_colors)

电流输出：

[((9, 9, 9), 1062),
 ((10, 10, 10), 700),
 ((8, 8, 8), 668),
 ((9, 7, 8), 586),
 ((9, 7, 10), 579)]

有没有真正的 Numpy 方法来解决这个问题？

标签： pythonnumpyimage-processingcolorsnumpy-slicing

Answer 1

方法#1

我们可以np.unique(.., axis=0, return_counts=True)用来获取每种唯一颜色的计数，然后np.argpartition获取其中前 N 个颜色的计数，以获得紧凑的矢量化解决方案 -

def topN_colors(img, N):
    unqc,C = np.unique(img.reshape(-1,img.shape[-1]), axis=0, return_counts=True)
    topNidx = np.argpartition(C,-N)[-N:]
    return unqc[topNidx], C[topNidx]

方法#2

另一个主要基于 24 位整数 2D 缩减，希望更有效的解决方案 -

# https://stackoverflow.com/a/57236217/ @tstanisl
def scalarize(x):
    # compute x[...,2]*65536+x[...,1]*256+x[...,0] in efficient way
    y = x[...,2].astype('u4')
    y <<= 8
    y +=x[...,1]
    y <<= 8
    y += x[...,0]
    return y

def topN_colors_v2(img, N):
    img2D = scalarize(img)
    unq,idx,C = np.unique(img2D, return_index=True, return_counts=True)
    topNidx = np.argpartition(C,-N)[-N:]
    return img.reshape(-1,img.shape[-1])[idx[topNidx]], C[topNidx]

请注意，argpartition不保留订单。为了保持秩序，使用range()它。更多信息。因此，np.argpartition取而代之-N的range(-N,0)是按升序获取颜色及其计数。对于降序，只需翻转最终输出。

样品验证

# Sample setup
np.random.seed(0)
# some random set colors
colors = np.array([[2,5,6],[1,2,3],[6,7,8],[5,3,1],[7,4,2]])

# Random image with colors chosen off colors
idx = np.random.randint(0,len(colors),(50,40))
img = colors[idx]
img = img.astype(np.uint8)

# Given that we know the unique colors as `colors` and the indices
# use to get the image `img, let's "manually" compute the 
# top N=2 colors and their counts
In [45]: count = np.bincount(idx.ravel())

In [46]: colors[count.argsort()[-2:]]
Out[46]: 
array([[1, 2, 3],
       [5, 3, 1]], dtype=uint8)

In [47]: count[count.argsort()[-2:]]
Out[47]: array([393, 446])

# Verify our solution
In [48]: topN_colors(img, N=2)
Out[48]: 
(array([[1, 2, 3],
        [5, 3, 1]], dtype=uint8),
 array([393, 446]))