A PyMaxflow example

maxflow_test.py

import numpy as np
import maxflow

from matplotlib import pyplot as plt

import networkx as nx

def plot_graph(nxgraph):
    X, Y = np.mgrid[:5, :5]
    aux = np.array([Y.ravel(), X[::-1].ravel()]).T
    positions = {i:aux[i] for i in xrange(25)}
    positions['s'] = (-1, 2)
    positions['t'] = (5, 2)
    
    # nxgraph.remove_nodes_from(['s', 't'])
    edge_labels = {}
    for u, v, d in nxgraph.edges(data=True):
        edge_labels[(u,v)] = d['weight']
    plt.clf()
    nx.draw(nxgraph, pos=positions)
    nx.draw_networkx_edge_labels(nxgraph, pos=positions, edge_labels=edge_labels, label_pos=0.3, font_size=7)
    plt.axis('equal')
    plt.show()

def create_graph():
    g = maxflow.Graph[float]()
    nodeids = g.add_grid_nodes((5,5))
    
    structure = np.array([[np.inf, 0, 0],
                          [np.inf, 0, 0],
                          [np.inf, 0, 0]
                         ])
    g.add_grid_edges(nodeids, structure=structure, symmetric=False)
    
    weights = np.array([[100, 110, 120, 130, 140]]).T + np.array([0, 2, 4, 6, 8])
    structure = np.zeros((3,3))
    structure[1,2] = 1
    g.add_grid_edges(nodeids, structure=structure, weights=weights, symmetric=False)
    
    structure = np.zeros((3,3))
    structure[0,1] = 1
    g.add_grid_edges(nodeids, structure=structure, weights=weights+100, symmetric=False)
    
    structure = np.zeros((3,3))
    structure[2,1] = 1
    g.add_grid_edges(nodeids, structure=structure, weights=weights+200, symmetric=False)
    
    left = nodeids[:, 0]
    g.add_grid_tedges(left, np.inf, 0)
    right = nodeids[:, -1]
    g.add_grid_tedges(right, 0, np.inf)
    
    return nodeids, g

if __name__ == '__main__':
    nodeids, g = create_graph()
    
    nxg = g.get_nx_graph()
    plot_graph(nxg)
    
    g.maxflow()
    print g.get_grid_segments(nodeids)