百度图片中有一个按照颜色搜图片的功能,其核心算法是提取图片主体颜色法,本文将使用python实现提取图片主体颜色算法。
百度将图片主体颜色归类到如下11中颜色:
颜色代码:
label_colors = [
[50, 50, 50],
[250, 250, 250],
[215, 68, 186],
[0, 101, 254],
[222, 32, 32],
[254, 191, 0],
[137, 43, 207],
[89, 167, 37],
[6, 183, 200],
[254, 108, 0],
[115, 52, 19],
]
因为cv2和PIL在读取gif图片时候会出现问题,所以我们使用imageio读取gif:
try:
tmp = imageio.mimread(img_path)
if tmp is not None:
imt = np.array(tmp)
imt = imt[0]
img = imt[:, :, 0:3]
img = Image.fromarray(img)
imageio_success = True
except:
imageio_success = False
if not imageio_success:
img = Image.open(img_path)
可以设置成截取图片中心部分:
input_size = 100
crop_size = 100
img = img[int((input_size - crop_size) / 2):int((input_size + crop_size) / 2), int((input_size - crop_size) / 2):int((input_size + crop_size) / 2), :]
颜色值量化,将256种颜色量化到8:
def quantization(img):
img = np.right_shift(img, quantization_factor)
img = np.left_shift(img, quantization_factor)
return img
提取主体颜色并排序:
def sort_color(img):
img_str = np.asarray(img, np.str)
img_str = img_str.reshape((-1, 3))
r_g_b = np.char.array(img_str[..., 0]) + '-' + np.char.array(img_str[..., 1]) + '-' + np.char.array(img_str[..., 2])
r_g_b = r_g_b.tolist()
r_g_b_count = Counter(r_g_b)
r_g_b_count = sorted(r_g_b_count.items(), key=lambda x: x[1], reverse=True)
max_count_color = r_g_b_count[0][0].split('-')
max_count_color = np.asarray(max_count_color, np.int)
return max_count_color, r_g_b_count[0][1]
从11种颜色中选取图片代表颜色:
def get_nearest_color(max_count_color):
nearest_index_euc = -1
nearest_distance_euc = 99999
nearest_index_abs = -1
nearest_distance_abs = 99999
for i, label_color in enumerate(label_colors):
euc = get_euc_distance(label_color, np.asarray(max_count_color))
if euc < nearest_distance_euc:
nearest_distance_euc = euc
nearest_index_euc = i
abs = get_abs_distance(label_color, np.asarray(max_count_color))
if abs < nearest_distance_abs:
nearest_distance_abs = abs
nearest_index_abs = i
return nearest_distance_euc, nearest_index_euc, nearest_distance_abs, nearest_index_abs
def get_euc_distance(color1, color2):
return np.linalg.norm(color1 - color2)
def get_abs_distance(color1, color2):
return np.sum(np.abs(color1 - color2))
返回值说明:
max_count_color 图片主体颜色
max_count_color_ratio 主题颜色占图片面积的比例
nearest_distance_euc 使用欧式距离,计算图片代表颜色,返回具体距离数值
nearest_index_euc 使用欧式距离计算出的图片代表颜色
nearest_distance_abs 使用L1距离,计算图片代表颜色,返回具体距离数值
nearest_index_abs 使用L1距离计算出的图片代表颜色
全部代码:
# coding=utf-8
# ================================================================
#
# File name : color_service.py
# Author : Faye
# E-mail : [email protected]
# Created date: 2022/12/10 22:45
# Description :
#
# ================================================================
import numpy as np
from collections import Counter
from PIL import Image
import imageio
label_colors = [
[50, 50, 50],
[250, 250, 250],
[215, 68, 186],
[0, 101, 254],
[222, 32, 32],
[254, 191, 0],
[137, 43, 207],
[89, 167, 37],
[6, 183, 200],
[254, 108, 0],
[115, 52, 19],
]
input_size = 100
crop_size = 100
quantization_factor = 5
def get_img_main_color(img):
"""
提取图片主体颜色
:param img: PIL格式的图片
:return:
max_count_color 图片主体颜色
max_count_color_ratio 主题颜色占图片面积的比例
nearest_distance_euc 使用欧式距离,计算图片代表颜色,返回具体距离数值
nearest_index_euc 使用欧式距离计算出的图片代表颜色
nearest_distance_abs 使用L1距离,计算图片代表颜色,返回具体距离数值
nearest_index_abs 使用L1距离计算出的图片代表颜色
"""
img = img.resize((input_size, input_size), Image.ANTIALIAS)
img = np.asarray(img)
if len(img.shape) == 2:
img = np.repeat(img[..., np.newaxis], 3, axis=-1)
if img.shape[-1] == 4:
img = img[..., 0:3]
img = img[int((input_size - crop_size) / 2):int((input_size + crop_size) / 2), int((input_size - crop_size) / 2):int((input_size + crop_size) / 2), :]
# is_gray = check_gray(img)
h, w, _ = img.shape
img = quantization(img)
max_count_color, max_count_color_ratio = sort_color(img)
nearest_color = get_nearest_color(max_count_color)
res = []
res.append(max_count_color)
res.append(max_count_color_ratio)
for n in nearest_color:
res.append(n)
return res
def quantization(img):
"""
颜色值量化
:param img:
:return:
"""
img = np.right_shift(img, quantization_factor)
img = np.left_shift(img, quantization_factor)
return img
def sort_color(img):
"""
提取主体颜色并排序
:param img:
:return:
"""
img_str = np.asarray(img, np.str)
img_str = img_str.reshape((-1, 3))
r_g_b = np.char.array(img_str[..., 0]) + '-' + np.char.array(img_str[..., 1]) + '-' + np.char.array(img_str[..., 2])
r_g_b = r_g_b.tolist()
r_g_b_count = Counter(r_g_b)
r_g_b_count = sorted(r_g_b_count.items(), key=lambda x: x[1], reverse=True)
max_count_color = r_g_b_count[0][0].split('-')
max_count_color = np.asarray(max_count_color, np.int)
return max_count_color, r_g_b_count[0][1]
def get_nearest_color(max_count_color):
"""
从11种颜色中选取图片代表颜色
:param max_count_color:
:return:
"""
nearest_index_euc = -1
nearest_distance_euc = 99999
nearest_index_abs = -1
nearest_distance_abs = 99999
for i, label_color in enumerate(label_colors):
euc = get_euc_distance(label_color, np.asarray(max_count_color))
if euc < nearest_distance_euc:
nearest_distance_euc = euc
nearest_index_euc = i
abs = get_abs_distance(label_color, np.asarray(max_count_color))
if abs < nearest_distance_abs:
nearest_distance_abs = abs
nearest_index_abs = i
return nearest_distance_euc, nearest_index_euc, nearest_distance_abs, nearest_index_abs
def get_euc_distance(color1, color2):
return np.linalg.norm(color1 - color2)
def get_abs_distance(color1, color2):
return np.sum(np.abs(color1 - color2))
def check_gray(img):
if len(img.shape) == 2:
return True
if img.shape[-1] == 4:
img = img[..., 0:3]
avg = np.mean(img, axis=-1)
r_diff = np.abs(img[..., 0] - avg)
g_diff = np.abs(img[..., 1] - avg)
b_diff = np.abs(img[..., 2] - avg)
diff_count = r_diff + g_diff + b_diff
diff_count_gt = np.where(diff_count > 10)[0]
if len(diff_count_gt) == 0:
return True
gray_per = len(diff_count_gt) / input_size**2
return gray_per > 0.03
if __name__ == '__main__':
img_path = r'1.jpg'
imageio_success = False
try:
tmp = imageio.mimread(img_path)
if tmp is not None:
imt = np.array(tmp)
imt = imt[0]
img = imt[:, :, 0:3]
img = Image.fromarray(img)
imageio_success = True
except:
imageio_success = False
if not imageio_success:
img = Image.open(img_path)
max_count_color, max_count_color_ratio, nearest_distance_euc, nearest_index_euc, nearest_distance_abs, nearest_index_abs = get_img_main_color(img)
print(max_count_color, max_count_color_ratio, nearest_distance_euc, nearest_index_euc, nearest_distance_abs, nearest_index_abs)