def generate_walklets_embeddings(input_file, k):
"""
Generates Walklets embeddings
Author: Alexander Ulyanov
:param input_file: the graph edge list file
:param k: the scale of the network
:return: the name of the file where embeddings are saved
"""
# Read the graph
G = read_graph(input_file, delimiter=r"\s+")
# Prepare the output file
output_folder = "output/embeddings/"
output_filename = input_file.split("/")[-1].replace(
".edgelist", ".walklets")
output_file = output_folder + output_filename
# Create and configure Walklets instance
walklets = Walklets(G, output_file, max_level=k, dimensions=128)
# Perform random walks with pre-defined skip factor
walks = walklets.create_walks()
# Generate embeddings
embeddings = walklets.create_embeddings(walks)
# Save embeddings
walklets.save_model(embeddings)
# Return the filename where the embeddings are saved
return output_file
def get_gc(self, region):
if region in self.gcs:
return self.gcs[region]
fname = os.path.join(self.graph_dir, region + '.graph')
g = graph_helper.read_graph(fname)
gc = graph_helper.GraphContainer(g)
self.gcs[region] = gc
return gc
def worker(f):
print(f)
name = f.split('.')[0]
g = graph_helper.read_graph(os.path.join(args.graph_dir, f))
g = g.clear_self()
wkt = g.convert_rs_to_wkt()
all_data = []
for linestring in wkt:
all_data.append(("AOI_0_{}_img0".format(name), linestring))
df = pd.DataFrame(all_data, columns=['ImageId', 'WKT_Pix'])
df.to_csv(os.path.join(args.save_dir, name + '.csv'), index=False)
def worker(region):
graph_path = os.path.join(args.graph_dir, region.name + ".graph")
if not os.path.isfile(graph_path):
print("graph: {} not found.".format(region.name))
return
g = graph_helper.read_graph(graph_path)
img_mask = np.zeros((args.img_size, args.img_size), dtype=np.uint8)
print(" |-> Generating {}_{}_{}".format(region.name, region.radius_y,
region.radius_x))
bound_pnt_min = geom.Point(region.radius_x,
region.radius_y).scale(args.img_size // 2)
for edge in g.edges.values():
src = edge.src(g).point.sub(bound_pnt_min)
dst = edge.dst(g).point.sub(bound_pnt_min)
cv.line(img_mask, (src.x, src.y), (dst.x, dst.y), 255, args.thickness)
# Draw mask
cv.imwrite(os.path.join(args.save_dir, region.name + ".png"), img_mask)
def main():
test_regions = get_regions(cfg.DIR.TEST_REGION_PATH)
if cfg.TEST.SINGLE_REGION != "":
test_regions = {
cfg.TEST.SINGLE_REGION: test_regions[cfg.TEST.SINGLE_REGION]
}
net = prepare_net().eval()
junction_nms_res = dict()
road_seg_filter_dict = dict()
graph_dict = dict()
for region_name in test_regions.keys():
graph_dict[region_name] = None
if cfg.TEST.INFER_STEP == "start":
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junction"),
exist_ok=True)
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT, "road"),
exist_ok=True)
print("infer segmentation start, INPUT_SIZE:{}".format(
cfg.TEST.CROP_SZ))
road_map_dict, junc_map_dict = infer_segmentation(
net, list(test_regions.keys()))
print("infer segmentation done")
print("junction nms start")
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junc_nms"),
exist_ok=True)
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junc_nms_vis"),
exist_ok=True)
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"road_seg_region_filter"),
exist_ok=True)
pool = Pool(cfg.TEST.CPU_WORKER)
if cfg.TEST.START_FROM_JUNC_PEAK:
for region_name in test_regions.keys():
junction_nms_res[region_name] = pool.apply_async(
junction_nms,
args=(
region_name,
junc_map_dict[region_name],
))
for region_name in test_regions.keys():
junction_nms_res[region_name] = junction_nms_res[
region_name].get()
del junc_map_dict
if cfg.TEST.START_FROM_ROAD_PEAK:
for region_name in test_regions.keys():
road_seg_filter_dict[region_name] = pool.apply_async(
road_seg_region_filter,
args=(
region_name,
road_map_dict[region_name],
))
for region_name in test_regions.keys():
road_seg_filter_dict[region_name] = road_seg_filter_dict[
region_name].get()
del road_map_dict
pool.close()
pool.join()
print("junction nms done")
elif cfg.TEST.INFER_STEP == "after_seg":
print("junction nms start")
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junc_nms"),
exist_ok=True)
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junc_nms_vis"),
exist_ok=True)
os.makedirs(os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"road_seg_region_filter"),
exist_ok=True)
pool = Pool(cfg.TEST.CPU_WORKER)
if cfg.TEST.START_FROM_JUNC_PEAK:
for region_name in test_regions.keys():
junc_nms_map = cv.imread(
os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"junction", region_name + '.png'), 0)
junction_nms_res[region_name] = pool.apply_async(
junction_nms, args=(
region_name,
junc_nms_map,
))
for region_name in test_regions.keys():
junction_nms_res[region_name] = junction_nms_res[
region_name].get()
if cfg.TEST.START_FROM_ROAD_PEAK:
for region_name in test_regions.keys():
road_seg_map = cv.imread(
os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT, "road",
region_name + '.png'), 0) / 255.
road_seg_filter_dict[region_name] = pool.apply_async(
road_seg_region_filter, args=(
region_name,
road_seg_map,
))
for region_name in test_regions.keys():
road_seg_filter_dict[region_name] = road_seg_filter_dict[
region_name].get()
pool.close()
pool.join()
print("junction nms end")
elif cfg.TEST.INFER_STEP in ["after_junc_nms", "given_junc_nms"
] and cfg.TEST.START_FROM_ROAD_PEAK:
for region_name in test_regions.keys():
road_seg_filter_dict[region_name] = cv.imread(
os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"road_seg_region_filter", region_name + ".png"),
0) / 255.
elif cfg.TEST.INFER_STEP == "after_graph_from_junc" and cfg.TEST.START_FROM_ROAD_PEAK:
for region_name in test_regions.keys():
road_seg_filter_dict[region_name] = cv.imread(
os.path.join(cfg.DIR.SAVE_SEG_DIR, cfg.TEST.CKPT,
"road_seg_region_filter", region_name + ".png"),
0) / 255.
if cfg.TEST.START_FROM_JUNC_PEAK:
graph_dict[region_name] = graph_helper.read_graph(
os.path.join(
cfg.DIR.SAVE_GRAPH_DIR,
'{}_{}'.format(cfg.TEST.CKPT, cfg.TEST.NUM_TARGETS),
'graphs_junc', '{}.graph'.format(region_name)))
img_cache = tileloader.TileCache(tile_dir=cfg.DIR.IMAGERY_DIR,
tile_size=cfg.TRAIN.IMG_SZ,
window_size=cfg.TEST.WINDOW_SIZE,
limit=cfg.TRAIN.PARALLEL_TILES)
paths = []
region_lst = list(test_regions.keys())
if not cfg.TEST.INFER_STEP == "after_graph_from_junc" and cfg.TEST.START_FROM_JUNC_PEAK:
for i, region_name in enumerate(region_lst):
tile_data = get_tile_data(test_regions[region_name],
img_cache,
junction_nms_res,
get_starting_locations=True)
paths.append(
model_utils.Path(i,
training=False,
gc=None,
tile_data=tile_data,
graph=None,
road_seg=None))
save_graph_dir = os.path.join(
cfg.DIR.SAVE_GRAPH_DIR,
'{}_{}'.format(cfg.TEST.CKPT, cfg.TEST.NUM_TARGETS), 'graphs_junc')
os.makedirs(save_graph_dir, exist_ok=True)
try:
iters, graph_dict = infer_anchor(
paths,
net,
region_lst=region_lst,
save_graph_dir=save_graph_dir,
batch_size=cfg.TEST.BATCH_SIZE_ANCHOR)
print(iters)
except:
for path in paths:
path.graph.save(
os.path.join(
save_graph_dir,
'except_{}.graph'.format(region_lst[path.idx])))
print(" Except save graph {}".format(region_lst[path.idx]))
if cfg.TEST.START_FROM_ROAD_PEAK:
if len(paths) == 0:
for i, region_name in enumerate(region_lst):
tile_data = get_tile_data(test_regions[region_name],
img_cache,
junction_nms_res,
get_starting_locations=False)
path = model_utils.Path(
i,
training=False,
gc=None,
tile_data=tile_data,
graph=graph_dict[region_name],
road_seg=np.ascontiguousarray(
road_seg_filter_dict[region_name].swapaxes(0, 1)))
paths.append(path)
else:
for i, region_name in enumerate(region_lst):
paths[i].road_seg = np.ascontiguousarray(
road_seg_filter_dict[region_name].swapaxes(0, 1))
paths[i].remove_graph_from_road_seg()
if cfg.TEST.START_FROM_JUNC_PEAK:
save_graph_dir = os.path.join(
cfg.DIR.SAVE_GRAPH_DIR, '{}_{}'.format(cfg.TEST.CKPT,
cfg.TEST.NUM_TARGETS),
'graphs_junc_road')
else:
save_graph_dir = os.path.join(
cfg.DIR.SAVE_GRAPH_DIR, '{}_{}'.format(cfg.TEST.CKPT,
cfg.TEST.NUM_TARGETS),
'graphs_road')
os.makedirs(save_graph_dir, exist_ok=True)
try:
iters, graph_dict = infer_anchor(
paths,
net,
region_lst=region_lst,
save_graph_dir=save_graph_dir,
batch_size=cfg.TEST.BATCH_SIZE_ANCHOR)
print(iters)
except:
for path in paths:
path.graph.save(
os.path.join(
save_graph_dir,
'except_{}.graph'.format(region_lst[path.idx])))
print(" Except save graph {}".format(region_lst[path.idx]))
post_process_graph(graph_dict)