Efficient Graph-Based Image Segmentation【python实现】

 Selective Search寻找候选区域的基础，尽管已经用numpy加速过了，但python实现下来效率着实低，用来学习算法，图个乐子。真要使用直接调用

from skimage.segmentation import felzenszwalb

from PIL import Image
import numpy as np


def gaussian_filter(img, K_size=3, sigma=1.3):
    if len(img.shape) == 3:
        H, W, C = img.shape
    else:
        img = np.expand_dims(img, axis=-1)
        H, W, C = img.shape

    ## Zero padding
    pad = K_size // 2
    out = np.zeros((H + pad * 2, W + pad * 2, C), dtype=np.float64)
    out[pad: pad + H, pad: pad + W] = img.copy().astype(np.float64)

    ## prepare Kernel
    K = np.zeros((K_size, K_size), dtype=np.float64)
    for x in range(-pad, -pad + K_size):
        for y in range(-pad, -pad + K_size):
            K[y + pad, x + pad] = np.exp(-(x ** 2 + y ** 2) / (2 * (sigma ** 2)))
    K /= (2 * np.pi * sigma * sigma)
    K /= K.sum()
    tmp = out.copy()
    # filtering
    for y in range(H):
        for x in range(W):
            for c in range(C):
                out[pad + y, pad + x, c] = np.sum(K * tmp[y: y + K_size, x: x + K_size, c])
    out = np.clip(out, 0, 255)
    return out


def graph_based_image_segmentation(img_src, k=50):
    g = []
    v = []
    e = []
    image = Image.open(img_src)
    # image.show()
    img_array = np.array(image, dtype=np.float64)
    img_array = gaussian_filter(img_array)
    a_offset = (img_array.shape[0] - 1) * (img_array.shape[1] - 1)
    b_offset = (img_array.shape[0] - 1) * img_array.shape[1] + a_offset

    def edgeToPoint(index):
        if i < a_offset:
            a = index // (img_array.shape[1] - 1) * img_array.shape[1] + index % (img_array.shape[1] - 1)
            b = a + img_array.shape[1] + 1
        elif a_offset <= index < b_offset:
            a = index - a_offset
            b = a + img_array.shape[1]
        else:
            a = (index - b_offset) // (img_array.shape[1] - 1) * img_array.shape[1] + ((index - b_offset) % (
                    img_array.shape[1] - 1))
            b = a + 1
        return int(a), int(b)

    img_edge = np.concatenate((np.sqrt(np.sum((img_array[:-1, :-1] - img_array[1:, 1:]) ** 2, 2)).reshape((-1, 1)),
                               np.sqrt(np.sum((img_array[:-1] - img_array[1:]) ** 2, 2)).reshape((-1, 1)),
                               np.sqrt(np.sum((img_array[:, :-1] - img_array[:, 1:]) ** 2, 2)).reshape((-1, 1))))
    img_edge = np.concatenate((np.arange(img_edge.shape[0]).reshape((-1, 1)), img_edge), axis=1)
    img_edge = img_edge[np.argsort(img_edge[:, 1])]
    v = np.arange(img_array.shape[0] * img_array.shape[1]).astype(np.int64)
    g = [np.array([0, i]) for i in range(img_array.shape[0] * img_array.shape[1])]
    for i, similarity in img_edge:
        point_a, point_b = edgeToPoint(i)
        g_i_1 = v[point_a]
        g_i_2 = v[point_b]
        if g_i_1 != g_i_2:
            if max(g[g_i_1][0] + k / len(g[g_i_1][1:]), g[g_i_2][0] + k / len(g[g_i_2][1:])) >= similarity:
                g[g_i_1] = np.append(g[g_i_1], g[g_i_2][1:], axis=0)
                g[g_i_1][0] = similarity
                v[g[g_i_2][1:]] = g_i_1
                g[g_i_2] = None
    new_img = np.zeros_like(img_array).reshape((-1, 3))
    for sub in g:
        if sub is not None:
            new_img[sub[1:]] = np.random.randint(0, 255, 3)
    new_img = new_img.reshape((img_array.shape[0], img_array.shape[1], 3))
    img = Image.fromarray(new_img.astype('uint8')).convert('RGB')
    img.show()


if __name__ == "__main__":
    graph_based_image_segmentation('R.jpg', 70)