百度图片中有一个按照颜色搜图片的功能,其核心算法是提取图片主体颜色法,本文将使用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 : xiansheng14@sina.com # 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 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)