Python图像处理之线路提取

#导入所需库和模块
import numpy as np
import math
import cv2
from sklearn.cluster import KMeans
from skimage.morphology import skeletonize
from scipy.signal import convolve2d
from simplification.cutil import simplify_coords
import webcolors
import matplotlib.pyplot as plt
%matplotlib inline

图像上传

img = cv2.imread('data/baby_yoda.jpeg')
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img_rgb)
plt.xticks([])
plt.yticks([])
plt.show();

边缘探测(edges detection)

def make_white_background(image, mask):
    """
    auxiliary function to make white background for an image
    :param image: color image
    :param mask: mask of 0 (background) and 1 (detected edges)
    :return: masked color image
    """
    result_image = image.copy()
    for idx_row in range(image.shape[0]):
        for idx_col in range(image.shape[1]):
            if mask[idx_row, idx_col] == 0:
                result_image[idx_row][idx_col] = [255, 255, 255]
    result_image = result_image.astype('uint8')
    return result_image

def adaptive_thresholding_edge_detection(image):
    """
    function for image thresholding with median filter and adaptive threshold
    :param image: source image in RGB
    :return: image of the same size in RGB format with white background and detected color edges
    """
    imgray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    img = cv2.medianBlur(imgray, 5)
    thresh = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 11, 2)

    mask = (1 - (thresh / 255)).astype('uint8')
    result_image = make_white_background(image, mask)

    return result_image, mask

def simple_thresholding_edge_detection(image):
    """
    function for image thresholding with simple binary threshold
    :param image: source image in RGB
    :return: image of the same size in RGB format with white background and detected color edges
    """
    imgray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    ret, thresh = cv2.threshold(imgray, 127, 255, 0)

    mask = 1 - (thresh / 255).astype('uint8')
    result_image = make_white_background(image, mask)

    return result_image, mask

边缘检测算法

def canny_edge_detection(image):
    """
    function for canny edge detection
    :param image: source image in RGB
    :return: image of the same size in RGB format with white background and detected color edges
    """
    imgray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    img = cv2.medianBlur(imgray, 5)
    thresh = cv2.Canny(img, 75, 150)

    mask = (thresh / 255).astype('uint8')
    result_image = make_white_background(image, mask)

    return result_image, mask

简单的阈值处理

img_rgb_simple_thrsh, mask_simple_thrsh = simple_thresholding_edge_detection(img_rgb)

plt.imshow(img_rgb_simple_thrsh)
plt.xticks([])
plt.yticks([])
plt.plot();

自适应阈值化

img_rgb_adaptive_thrsh, mask_adaptive_thrsh = adaptive_thresholding_edge_detection(img_rgb)

plt.imshow(img_rgb_adaptive_thrsh)
plt.xticks([])
plt.yticks([])
plt.plot();

边缘检测

img_rgb_canny, mask_canny = canny_edge_detection(img_rgb)

plt.imshow(img_rgb_canny)
plt.xticks([])
plt.yticks([])
plt.plot();

 

矢量化(线选择)

支持功能

def distance(x1, y1, x2, y2):
    """
    function for calculating distance between two points
    :param x1: x coordinate of the first point
    :param y1: y coordinate of the first point
    :param x2: x coordinate of the second point
    :param y2: y coordinate of the second point
    :return: distance value
    """
    return math.sqrt(math.pow(x1 - x2, 2) + math.pow(y1 - y2, 2))

def line_length(line):
    """
    function for calculating length of a line
    :param line: as sequence of points
    :return: length of a line
    """
    line_len = 0
    for p in range(0, len(line) - 1):
        (x1, y1) = (line[p][1], line[p][0])
        (x2, y2) = (line[p + 1][1], line[p + 1][0])
        line_len = line_len + distance(x1, y1, x2, y2)
    return line_len

def find_neighbour_points(img, x, y, threshold=1):
    """
    function for finding neighbour points
    :param img: tested image
    :param x: x coordinate
    :param y: y coordinate
    :param threshold: thresholding value
    :return: detected lines
    """
    out_points = []

    if x - 1 >= 0 and y - 1 >= 0:
        if img[x - 1][y - 1] == threshold:
            out_points.append((x - 1, y - 1))

    if x - 1 >= 0:
        if img[x - 1][y] == threshold:
            out_points.append((x - 1, y))

    if x - 1 >= 0 and y + 1 < img.shape[1]:
        if img[x - 1][y + 1] == threshold:
            out_points.append((x - 1, y + 1))

    if y - 1 >= 0:
        if img[x][y - 1] == threshold:
            out_points.append((x, y - 1))

    if y + 1 < img.shape[1]:
        if img[x][y + 1] == threshold:
            out_points.append((x, y + 1))

    if x + 1 < img.shape[0] and y - 1 >= 0:
        if img[x + 1][y - 1] == threshold:
            out_points.append((x + 1, y - 1))

    if x + 1 < img.shape[0]:
        if img[x + 1][y] == threshold:
            out_points.append((x + 1, y))

    if x + 1 < img.shape[0] and y + 1 < img.shape[1]:
        if img[x + 1][y + 1] == threshold:
            out_points.append((x + 1, y + 1))

    return out_points

def find_lines(skeleton):
    """
    function for finding continuous sequence which may be treated as lines
    :param skeleton: result of skimage skeletonize
    :return: detected lines
    """
    lines = []
    for i in range(0, skeleton.shape[0]):
        for j in range(0, skeleton.shape[1]):
            if skeleton[i][j] == 1:
                line = []
                line.append((i, j))
                neb = find_neighbour_points(skeleton, i, j)
                skeleton[i][j] = 0
                while len(neb) != 0:
                    for n in neb:
                        skeleton[n[0], n[1]] = 0
                    line.append(neb[0])
                    nn = find_neighbour_points(skeleton, neb[0][0], neb[0][1])
                    neb.pop(0)
                    for n in neb:
                        nn.append(n)
                    neb = [n for n in nn]
                line = list(map(lambda point: (point[1], point[0]), line))
                lines.append(line)
    return lines

 

 

#线条骨架化
skeleton = skeletonize(mask_adaptive_thrsh)

#删除线的交叉点
neighbour_matrix = convolve2d(skeleton, np.ones((3, 3)), mode='same')
skeleton[neighbour_matrix > 3] = 0

#线路矢量化
lines = find_lines(skeleton)

#去除环路
new_lines = []
for line in lines:
    lenlen = []
    p = 0
    while p < len(line) - 1:
        if distance(line[p][0], line[p][1], line[p + 1][0], line[p + 1][1]) > 2:
            lenlen.append(p)
        p = p + 1
    if len(lenlen) == 1:
        i = lenlen[0]
        new_lines.append(line[i + 1: len(line)])
        new_lines.append(line[0: i + 1])
    else:
        new_lines.append(line)

#清除垃圾线
eps_line_len = 1
new_lines = [line for line in new_lines if len(line) > eps_line_len]
new_lines = [line for line in new_lines if line_length(line) > 5]

用Ramer-Douglas-Pecker算法简化线条

使用算法的示例

simplified = simplify_coords(new_lines[0], 1.0)

简化后的结果

print('Исходные координаты первой линии')
print(str(new_lines[0]))

print('Упрощенные координаты первой линии')
print(str(simplified))

按颜色分类的线条聚类

支持性功能

def get_colors_of_line(image, line):
    """
    function to get color of line
    :param image: original RGB image
    :param line: detected line as list of points
    :return: color
    """
    return [image[elem[1]][elem[0]] for elem in line]

def get_colors(image, lines):
    """
    function to get colors of lines
    :param image: original RGB image
    :param lines: detected lines as list of points
    :return: colors
    """
    rgb_lines = []
    for line in lines:
        t = get_colors_of_line(image, line)
        (r, g, b) = (0, 0, 0)
        for tt in t:
            r = r + tt[0]
            g = g + tt[1]
            b = b + tt[2]
        (r, g, b) = (int(r / len(t)), int(g / len(t)), int(b / len(t)))
        rgb_lines.append((r, g, b))
    return rgb_lines

number_of_clusters = 16

rgb_lines = get_colors(img_rgb, new_lines)
rgb_lines_array = np.array(rgb_lines, np.uint8)
lines_array = np.array(new_lines)

clusterer = KMeans(n_clusters=number_of_clusters)
clusterer.fit(rgb_lines_array)
labels = clusterer.labels_
centers = clusterer.cluster_centers_

lines_by_clusters = [list(lines_array[labels == i]) for i in range(number_of_clusters)]
cluster_color = list(map(lambda color: tuple(map(int, color)), centers))

绘制结果

支持性功能

def draw_lines_by_color(lines, debug_image, color, line_width=3):
    """
    function for drawing lines
    :param lines: detected lines as list of points
    :param debug_image: debug image
    :param color: color
    :param line_width: width of a line
    :return:
    """
    for line in lines:
        for p0, p1 in zip(line[:-1], line[1:]):
            x0, y0 = p0
            x1, y1 = p1
            cv2.line(debug_image, (int(x0), int(y0)), (int(x1), int(y1)), color, line_width)

def get_color_name(rgb_triplet):
    """
    function to get closest color name by RGB triplet
    :param rgb_triplet: (R, G, B)
    :return: identified color as str
    """
    min_colors = {}
    for name, hex_value in webcolors.CSS3_NAMES_TO_HEX.items():
        r_c, g_c, b_c = webcolors.hex_to_rgb(hex_value)
        rd = (r_c - rgb_triplet[0]) ** 2
        gd = (g_c - rgb_triplet[1]) ** 2
        bd = (b_c - rgb_triplet[2]) ** 2
        min_colors[(rd + gd + bd)] = name
    return min_colors[min(min_colors.keys())]

plt.figure(figsize=(15, 120))

for i, lines in enumerate(lines_by_clusters):
    plt.subplot(number_of_clusters, 2, 2 * i + 1)
    plt.imshow(img_rgb)
    plt.xticks([])
    plt.yticks([])
    
    plt.subplot(number_of_clusters, 2, 2 * i + 2)
    debug_image = (255 + np.zeros((img_rgb.shape[0], img_rgb.shape[1], 3))).astype('uint8')
    draw_lines_by_color(lines, debug_image, cluster_color[i], line_width=3)
    plt.imshow(debug_image.astype('uint8'))
    plt.title(get_color_name(cluster_color[i]))
    plt.xticks([])
    plt.yticks([])

plt.plot();