百度图片源码流出~按照颜色搜图片~提取图片主体颜色

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

你可能感兴趣的:(python,人工智能,计算机视觉,opencv)