使用Python查找具有相似调色板的图像

Question

假设图库中有 10,000 张 JPEG、PNG 图像，如何找到与所选图像具有相似调色板的所有图像并按相似度降序排序？

Answer 1

为每个图像构建颜色直方图。然后，当您想要将图像与集合相匹配时，只需按照其直方图与所选图像的直方图的接近程度对列表进行排序即可。

桶的数量取决于您想要的准确度。组合成存储桶的数据类型将定义您如何确定搜索的优先级。

例如，如果您对色调最感兴趣，那么您可以定义图像的每个单独像素属于哪个存储桶：

def bucket_from_pixel(r, g, b):
    hue = hue_from_rgb(r, g, b) # [0, 360)
    return (hue * NUM_BUCKETS) / 360

如果您还想要通用匹配器，那么您可以根据完整的 RGB 值来选择桶。

使用PIL，可以使用内置的

histogram

函数。可以使用您想要的任何距离度量来计算“接近度”直方图。例如，L1 距离可以是：

hist_sel = normalize(sel.histogram())
hist = normalize(o.histogram()) # These normalized histograms should be stored

dist = sum([abs(x) for x in (hist_sel - hist)])

L2 是：

dist = sqrt(sum([x*x for x in (hist_sel - hist)]))

Normalize

只是强制直方图的总和等于某个常数值（1.0 就可以了）。这很重要，以便可以正确地将大图像与小图像进行比较。如果您要使用 L1 距离，那么您应该在

normalize

中使用 L1 度量。如果是 L2，则 L2。

Answer 2

另一种解决方案是在每个调色板上运行 Kmeans 对颜色进行分组，按色调对每个调色板进行排序，然后使用余弦相似度来查找最相似的图像。

这里有一段代码，可以在图像文件夹中找到与 ref_image 最相似的图像：

from PIL import Image
import os
import numpy as np
from sklearn.cluster import MiniBatchKMeans
from numpy.linalg import norm
from tqdm import tqdm
import colorsys


def extract_palette(image, num_colors):
  image_array = np.array(image)
  pixels = image_array.reshape((-1, 3))
  kmeans = MiniBatchKMeans(n_clusters=num_colors, n_init='auto')
  kmeans.fit(pixels)
  colors = kmeans.cluster_centers_
  return colors


def order_vector_by_hue(colors):
  hsv_colors = [colorsys.rgb_to_hsv(*color) for color in colors]
  ordered_hsv_colors = sorted(hsv_colors, key=lambda x: x[0])
  ordered_rgb_colors = [tuple(int(c * 255) for c in colorsys.hsv_to_rgb(*hsv)) for hsv in ordered_hsv_colors]
  return np.array(ordered_rgb_colors)


def cosine_sim(u, v):
  return np.dot(u, v) / (norm(u) * norm(v))


if __name__ == "__main__":
  ref_image_path = '<your-path>'
  folder = '<your-image-folder>'

  files = os.listdir(folder)

  print('processing ref image')
  image = Image.open(image_path)
  ref_colors = extract_palette(image, num_colors=32)
  ref_colors = order_vector_by_hue(colors).reshape(-1)

  print('processing candidate images')
  selected_image_path = None
  max_similarity = -1 # value for the most dissimilar possible image
  for image_path in files:
    image = Image.open(image_path)
    colors = extract_palette(image, num_colors=32)
    colors = order_vector_by_hue(colors).reshape(-1)
    similarity = cosine_sim(ref_colors, colors)
    if similarity > max_similarity:
      max_similarity = similarity
      selected_image_path = image_path

  print(f'most similar image: {selected_image_path}')