gt_segmentation.py

#!/usr/bin/env python3

# The main script used for segmentation with a functional algorithm.
# If needed, change: 
#   node_id_param_file location, 
#   orientation
#   edgelist file name and location
#   OUTPUT: a gt file (change name as needed)

import graph_tool
from graph_tool import draw
import time
import re
import json
import math
from os import path
import shapefile
import datetime
from numpy import loadtxt


##########################
# load the edgelist file #
##########################

print('Loading edgelist file...')
g = graph_tool.load_graph_from_csv('./edgelist_52.csv', skip_first=False, directed=True, hashed=True, eprop_names=['nl_cost_function'], csv_options={'delimiter': ' ', 'quotechar': '"'})
print('Loading done')


#####################
# Set property maps #
#####################

g.set_fast_edge_removal()

vp = g.new_vp('object') # this object will be a list
g.vp['segment'] = vp

pos_first = g.new_vp('object') 
g.vp['pos_first'] = pos_first

pos_last = g.new_vp('object') 
g.vp['pos_last'] = pos_last

f_av_r = g.new_vp('float') 
g.vp['first_av_rad'] = f_av_r

f_av_ros = g.new_vp('float') 
g.vp['first_av_ros'] = f_av_ros

l_av_r = g.new_vp('float') 
g.vp['last_av_rad'] = l_av_r

l_av_ros = g.new_vp('float') 
g.vp['last_av_ros'] = l_av_ros

nl_cf = g.new_ep('float')
g.ep['nl_cost_function'] = nl_cf

print('Starting processing')

##############
# Processing #
##############

counter = 1
verts_to_del = []

t2 = time.perf_counter()

for v in g.vertices():
    if (counter % 10000 == 0):
        t3 = time.perf_counter()
        print(f'Processing {counter} of {g.num_vertices()}')
        print(t3 - t2)
        t2 = time.perf_counter()

    counter += 1

    if not g.vp.segment[v]:
        g.vp.segment[v] = [int(g.vp.name[v])]

    #print (f'processing {g.vp.name[v]}...:')
    #print (f'In deg: {v.in_degree()}, out deg: {v.out_degree()}')

    if v.in_degree() == 1:
        v2 = next(v.in_edges()).source()
        
        if v2.out_degree() == 1:
            g.remove_edge(g.edge(v2, v))
            g.vp.segment[v2].append(int(g.vp.name[v]))

            for v3 in v.out_neighbors():
                g.remove_edge(g.edge(v, v3))
                g.vp.segment[v].append(int(g.vp.name[v3]))
                new_edge = g.add_edge(v2, v3)

    if(v.in_degree() == 0 and v.out_degree() == 0):
        verts_to_del.append(v)

    
print(f'Done processing ! New graph: {g}')
print(f'Empty nodes: {len(verts_to_del)}. Deleting empty nodes...')

t2 = time.perf_counter()
g.remove_vertex(verts_to_del, fast=True)
t3 = time.perf_counter()
print(t3 - t2)

print('Saving graph...')
g.save('global_segmented_anti_50.gt')
print('All done.')