色卡矩形检测

最近需要做图像的颜色校正，其中有一步需要把图像中对比用的色卡分离出来，然后才能求解回归方程。思想很简单，二值化后做直线检测，然后组成矩形。如果二值化后直接调API做矩形检测的话，因为是色卡，所以有些矩形是检测不出来的，所以用水平线和垂直线相交的办法来定位矩形，最后去除面积太小的即可。
先给个定妆照：

可以看到背景和需要的色块全部分开了，效果还不错
=================== 我是一切不给源码的都是耍流氓的分割线 ================

# -*- coding:utf-8 -*-
import cv2
import numpy as np


def _img_split_with_shadow(gray_img, threshold_value=180):
    """
    :param binary_img: 读入的灰度图
    :param img_show:
    :return: 水平和垂直线的坐标集合
    """
    h = gray_img.shape[0]
    w = gray_img.shape[1]

    # 按行求和
    sum_x = np.sum(gray_img, axis=1)
    # 按列求和
    sum_y = np.sum(gray_img, axis=0)

    h_line_index = np.argwhere(sum_x == 0)
    v_line_index = np.argwhere(sum_y == 0)

    h_line_index = np.reshape(h_line_index, (h_line_index.shape[0],))
    v_line_index = np.reshape(v_line_index, (v_line_index.shape[0],))

    h_line = []
    v_line = []

    for i in range(len(h_line_index) - 1):
        if h_line_index[i + 1] - h_line_index[i] > 2:
            h_line.append((0, h_line_index[i + 1], w - 1, h_line_index[i + 1]))
            h_line.append((0, h_line_index[i], w - 1, h_line_index[i]))

    for i in range(len(v_line_index) - 1):
        if v_line_index[i + 1] - v_line_index[i] > 2:
            v_line.append((v_line_index[i + 1], 0, v_line_index[i + 1], h - 1))
            v_line.append((v_line_index[i], 0, v_line_index[i], h - 1))

    return h_line, v_line


def _combine_rect(h_lines, v_lines):
    """
    :param h_lines: 平行直线集合
    :param v_lines: 垂直直线集合
    :return: 返回由 h_lines 和 v_lines 组成的矩形集合
    """
    rects = []

    x_axis = sorted(set([item[0] for item in v_lines]))
    y_axis = sorted(set([item[1] for item in h_lines]))

    point_list = []
    for y in y_axis:
        point = []
        for x in x_axis:
            point.append((y, x))
        point_list.append(point)

    for y_index in range(len(y_axis) - 1):
        for x_index in range(len(x_axis) - 1):
            area = abs((y_axis[y_index + 1] - y_axis[y_index]) * (x_axis[x_index + 1] - x_axis[x_index]))
            rects.append([(y_axis[y_index], x_axis[x_index],
                          y_axis[y_index + 1], x_axis[x_index + 1]), area])
    # 按面积降序排序
    rects.sort(key = lambda  ele: ele[1], reverse=True)
    areas = [ele[1] for ele in rects]

    # 找到相邻差值最大的序号
    max = -1
    index = 0
    for i in range(len(areas) - 1):
        dif = areas[i] - areas[i + 1]
        if max < dif:
            max = dif
            index = i + 1

    return [ele[0] for ele in rects[0:index]]


def img_split(img_file, img_show=False):
    """
    :param img_file: 输入图片路径
    :param img_show: 是否显示
    :return: 分割后的子图像rect列表
    """
    img = cv2.imread(img_file, 1)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 7)
    binary = cv2.bitwise_not(binary)
    h = img.shape[0]
    w = img.shape[1]
    rate = h // w if h > w else w // h

    h_line_shadow, v_line_shadow = _img_split_with_shadow(binary)
    h_line = h_line_shadow
    v_line = v_line_shadow
    rects = _combine_rect(h_line, v_line)
    split_imgs = []
    for rect in rects:
        split_imgs.append(img[rect[0]:rect[2], rect[1]:rect[3]])

    if img_show:
        for rect in rects:
            cv2.rectangle(img, (rect[1], rect[0]), (rect[3], rect[2]), (0, 255, 0), 2)
        # lines = h_line + v_line
        # for x1, y1, x2, y2 in lines:
        #     cv2.line(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
        img = cv2.resize(img, (int(h * 0.8), int(h * 0.8 / rate)))
        cv2.imshow('cece', img)
        cv2.waitKey()
    return split_imgs


if __name__ == '__main__':
    imgs = img_split('std_color.png', True)
    # for index, img in enumerate(imgs):
    #     cv2.imshow(str(index), img)
    # cv2.waitKey(0)