使用python和igraph從圖像構造knn圖

Question

我已經實現了從圖像構造圖形的方法。 該方法基於KNN。

基本上，每個像素代表一個頂點，並且必須將每個像素與k個最近的鄰居相連。

該腳本很簡單，但是非常慢。 我試圖優化歐幾里得距離的計算以及增加邊的步驟。

有人對我的代碼有什么建議嗎？

最慢的步驟是歐幾里得距離的計算。 由於計算所有頂點的所有頂點的距離，因此該距離為n^2 。 例如，尺寸為600x375的圖片具有225000個折頁。

執行：

python file.py -f image.jpg -k 10

碼：

import Image
import math
from optparse import OptionParser
import igraph

def euclidian_distance(x1,y1,r1,g1,b1,x2,y2,r2,g2,b2):
    return math.sqrt(
        (x1 - x2) ** 2 +
        (y1 - y2) ** 2 +
        (r1 - r2) ** 2 +
        (g1 - g2) ** 2 +
        (b1 - b2) ** 2
    )

def _plot_xy(g):
    visual_style = {}
    visual_style["vertex_shape"] = "circle"
    visual_style["label_color"] = "white"
    visual_style["edge_color"] = "black"
    visual_style["edge_width"] = 0.2
    visual_style["vertex_size"] = 0.5

    layout = []
    for vertex in g.vs():
        layout.append((vertex["x"],vertex["y"]))

    visual_style["layout"] = layout
    visual_style["bbox"] = (200, 200)
    visual_style["margin"] = 10
    igraph.plot(g, **visual_style)

if __name__ == '__main__':

    parser = OptionParser()

    usage = "usage: python %prog [options] args ..."
    description = """Description"""
    parser.add_option("-f", "--file", dest="filename", help="read FILE", metavar="FILE")
    parser.add_option("-k", "--knn", dest="k", help="knn")

    (options, args) = parser.parse_args()
    filename = options.filename
    k = int(options.k)

    if filename is None:
        parser.error("required -f [filename] arg.")

    g = igraph.Graph()
    im = Image.open(filename)
    pix = im.load()
    for j in range(0,im.size[1]):
        for i in range(0,im.size[0]):
            g.add_vertex()
            vertex = g.vs[g.vcount()-1]
            vertex["name"] = vertex.index
            vertex["x"] = i
            vertex["y"] = j
            vertex["r"] = pix[i,j][0]
            vertex["g"] = pix[i,j][1]
            vertex["b"] = pix[i,j][2]

    // --> This step is very slow
    for v in g.vs():
        set_distance = dict()
        for n in g.vs():
            distance = euclidian_distance(v["x"],v["y"],v["r"],v["g"],v["b"],n["x"],n["y"],n["r"],n["g"],n["b"])
            set_distance[n.index] = distance
        sorted_set_distance = sorted(set_distance.items(), key=lambda set_distance: set_distance[1])
        v["distance"] = sorted_set_distance[:k]

    edges = []
    weight = []
    for v in g.vs():
        for n in v["distance"]:
            edges += [(v.index, n[0])]
            weight.append(n[1])

    g.add_edges(edges)
    g.es["weight"] = weight

    _plot_xy(g)

    g.write(filename.split('.')[0]+".edgelist", format='ncol')

Answer 1

不用為所有節點對計算歐幾里得距離，而是從節點構建一個kd-tree ，然后簡單地使用kd-tree獲取最近的鄰居。 這將大大減少距離計算的時間。 SciPy包含kd-tree的高效實現 ，因此無需重新發明輪子。

Answer 2

根據Tamas的答案，我修改了原始代碼。 比原始代碼快：

import math
from optparse import OptionParser
import igraph
from scipy import spatial
import numpy as np

if __name__ == '__main__':

    parser = OptionParser()

    usage = "usage: python %prog [options] args ..."
    description = """Description"""
    parser.add_option("-f", "--file", dest="filename", help="read FILE", metavar="FILE")
    parser.add_option("-k", "--knn", dest="k", help="knn")

    (options, args) = parser.parse_args()
    filename = options.filename
    k = int(options.k)

    if filename is None:
        parser.error("required -f [filename] arg.")

    graph = igraph.Graph()
    im = Image.open(filename)
    pix = im.load()
    x, y, r, g, b = [], [], [], [], []
    for j in range(0,im.size[1]):
        for i in range(0,im.size[0]):
            graph.add_vertex()
            vertex = graph.vs[graph.vcount()-1]
            vertex["name"] = vertex.index
            vertex["x"] = i
            vertex["y"] = j
            vertex["r"] = pix[i,j][0]
            vertex["g"] = pix[i,j][1]
            vertex["b"] = pix[i,j][2]
            x.append(i)
            y.append(j)
            r.append(pix[i,j][0])
            g.append(pix[i,j][1])
            b.append(pix[i,j][2])

    x = np.array(x)
    y = np.array(y)
    r = np.array(r)
    g = np.array(g)
    b = np.array(b)
    tree = spatial.KDTree(zip(x.ravel(), y.ravel(), r.ravel(), g.ravel(), b.ravel()))

    edges = []
    weight = []
    for v in graph.vs():
        pts = np.array([[v["x"], v["y"], v["r"], v["g"], v["b"]]])
        list_nn = tree.query(pts, k=k);
        for idx, nn in enumerate(list_nn[1][0]):
            edges += [(v.index, nn)]
            weight.append(1/(1+list_nn[0][0][idx]))

    graph.add_edges(edges)
    graph.es["weight"] = weight

    graph.write(filename.split('.')[0]+".edgelist", format='ncol')