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

Question

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

首先，我有一个 PIL.Image 对象，我在 x,y 坐标上循环并将其计算在默认字典中。我已经用 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 方法来解决这个问题？

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()。 More info。因此，在 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]))