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

Find Top10(n) RGB colors in Numpy

我有一个吃几秒钟的功能。该函数应该 return 给定图像中的 Top(n) 颜色。 return 必须排序,所以我可以使用第一、第二、第三最顶部颜色的 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 方法来解决这个问题吗?

方法 #1

我们可以将np.unique(.., axis=0, return_counts=True) to get counts for each unique color and then 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 缩减以获得更有效的解决方案 -

#  @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]))