in libs/apls/skeletonize.py [0:0]
def build_wkt_dir(indir, outfile, out_gdir,
thresh=0.3,
im_prefix='',
debug=False,
add_small=True,
fix_borders=True,
img_shape=(1300, 1300),
skel_replicate=5, skel_clip=2,
img_mult=255,
hole_size=300, cv2_kernel_close=7, cv2_kernel_open=7,
min_subgraph_length_pix=50,
min_spur_length_pix=16,
spacenet_naming_convention=False,
num_classes=1,
max_out_size=(100000, 100000),
use_medial_axis=True,
skeleton_band='all',
kernel_blur=27,
min_background_frac=0.2,
n_threads=12,
verbose=False,
super_verbose=False):
'''Execute built_graph_wkt for an entire folder
Split image name on AOI, keep only name after AOI. This is necessary for
scoring'''
im_files = np.sort([z for z in os.listdir(indir) if z.endswith('.tif')])
nfiles = len(im_files)
if n_threads is not None:
n_threads = min(n_threads, nfiles)
params = []
for i, imfile in enumerate(im_files):
if verbose: print("\n", i+1, "/", nfiles, ":", imfile)
img_loc = os.path.join(indir, imfile)
if verbose: print(" img_loc:", img_loc)
if spacenet_naming_convention:
im_root = 'AOI' + imfile.split('AOI')[-1].split('.')[0]
else:
im_root = imfile.split('.')[0]
if len(im_prefix) > 0:
im_root = im_root.split(im_prefix)[-1]
if verbose: print(" im_root:", im_root)
out_gpickle = os.path.join(out_gdir, imfile.split('.')[0]+'.gpickle')
if verbose: print(" out_gpickle:", out_gpickle)
param_row = (img_loc, out_gpickle, thresh, \
debug, fix_borders, \
img_shape,\
skel_replicate, skel_clip, \
img_mult, hole_size, \
cv2_kernel_close, cv2_kernel_open,\
min_subgraph_length_pix,\
min_spur_length_pix,\
max_out_size,\
use_medial_axis,\
num_classes,\
skeleton_band, \
kernel_blur,\
min_background_frac,\
verbose)
params.append(param_row)
# Compute skeleton graph (no geospatial info yet)
if n_threads is None or n_threads > 1:
if n_threads is None:
print("Running in parallel using all threads ...")
else:
print("Running in parallel using {} threads ...".format(n_threads))
# with Pool(n_threads) as pool:
# tqdm(pool.map(img_to_ske_G, params), total=len(params))
# Replace python multiprocessing.Pool with p_tqdm:
# https://github.com/swansonk14/p_tqdm
p_umap(img_to_ske_G, params, num_cpus=n_threads)
else:
print("Running in sequential using 1 thread ...")
for param in tqdm(params):
img_to_ske_G(param)
# Build wkt_list from non-geo skeleton graph (single-threaded)
all_data = []
for gpickle in os.listdir(out_gdir):
gpath = os.path.join(out_gdir, gpickle)
imfile = gpickle.split('.')[0] + '.tif'
if spacenet_naming_convention:
im_root = 'AOI' + imfile.split('AOI')[-1].split('.')[0]
else:
im_root = imfile.split('.')[0]
if len(im_prefix) > 0:
im_root = im_root.split(im_prefix)[-1]
G = nx.read_gpickle(gpath)
wkt_list = G_to_wkt(G, add_small=add_small,
verbose=verbose, super_verbose=super_verbose)
# add to all_data
for v in wkt_list:
all_data.append((im_root, v))
# save to csv
df = pd.DataFrame(all_data, columns=['ImageId', 'WKT_Pix'])
df.to_csv(outfile, index=False)
return df