Spaces:

Egrt
/

HEAT

Runtime error

App Files Files Community

HEAT / s3d_preprocess /generate_planar_graph.py

Egrt

init

424188c over 1 year ago

raw history blame contribute delete

No virus

24.8 kB

	import os
	import json
	import cv2
	import numpy as np
	from collections import defaultdict
	from scipy import ndimage


	def generate_graph(annot, image_path, out_path):
	lines = annot['lines']
	junctions = annot['junctions']
	line_junc_mat = np.array(annot['lineJunctionMatrix'])
	planes = annot['planes']
	plane_line_mat = annot['planeLineMatrix']
	plane_to_line = np.array(annot['planeLineMatrix'])
	line_to_plane = plane_to_line.T
	semantics = annot['semantics']

	all_room_edges = get_room_edges(semantics, planes, lines, junctions, plane_to_line, line_junc_mat)

	all_room_edges = filter_rooms(all_room_edges, im_size=256)

	all_colinear_pairs = find_all_colinear_paris(all_room_edges)
	colinear_sets = combine_colinear_edges(all_colinear_pairs)

	for colinear_set in colinear_sets:
	edges_to_merge = list(colinear_set)
	edges_to_merge = sorted(edges_to_merge, key=lambda x: -x[0])
	merged_edges = merge_edges(edges_to_merge)

	for merged_edge, old_edge in zip(merged_edges, edges_to_merge):
	if len(merged_edge) > 0:
	assert merged_edge[0][0] == old_edge[0]
	room_idx = merged_edge[0][0]

	if len(merged_edge) == 1 and merged_edge[0] == old_edge:
	continue
	# change room graph accordingly
	replaced_idx = all_room_edges[room_idx].index(old_edge[1])
	all_room_edges[room_idx].pop(replaced_idx)
	for new_idx, new_edge in enumerate(merged_edge):
	insert_idx = new_idx + replaced_idx
	all_room_edges[room_idx].insert(insert_idx, new_edge[1])
	else:
	room_idx = old_edge[0]
	replaced_idx = all_room_edges[room_idx].index(old_edge[1])
	all_room_edges[room_idx].pop(replaced_idx)

	# take intersection for every rooms to recover the room structure
	refined_room_edges = [adjust_room_edges(room_edges) for room_edges in all_room_edges]

	# clean every room loop by removing I-shape corners
	cleaned_room_edges = clean_room_edges(refined_room_edges)

	global_graph = defaultdict(list)
	for room_edges in cleaned_room_edges:
	for edge in room_edges:
	c1, c2 = edge
	global_graph[c1] += [c2, ]
	global_graph[c2] += [c1, ]
	for corner in global_graph:
	global_graph[corner] = list(set(global_graph[corner]))

	annot_path = os.path.join(out_path, 'annot.npy')
	np.save(annot_path, global_graph)

	# draw the planar graph on the density map
	viz_image = cv2.imread(image_path)
	for c, connections in global_graph.items():
	for other_c in connections:
	cv2.line(viz_image, (int(c[0]), int(c[1])), (int(other_c[0]), int(other_c[1])), (255, 255, 0), 2)
	for c in global_graph.keys():
	cv2.circle(viz_image, (int(c[0]), int(c[1])), 3, (0, 0, 255), -1)

	# viz_image = np.zeros([256, 266, 3]).astype(np.uint8)
	# room_idx = 0
	# for room_edges in cleaned_room_edges:
	# for line_idx, edge in enumerate(room_edges):
	# c1, c2 = np.array(edge).astype(np.int)
	# cv2.line(viz_image, tuple(c1), tuple(c2), (255, 255, 0), 2)
	# cv2.circle(viz_image, tuple(c1), 3, (0, 0, 255), -1)
	# cv2.circle(viz_image, tuple(c2), 3, (0, 0, 255), -1)
	# room_idx += 1
	cv2.imwrite(os.path.join(out_path, 'planar_graph.png'), viz_image)


	def get_room_edges(semantics, planes, lines, junctions, plane_to_line, line_junc_mat):
	room_edges = list()
	for semantic in semantics:
	plane_ids = semantic['planeID']
	label = semantic['type']
	if label in ['door', 'window', 'outwall']: # skip non-room elements
	continue
	all_planes = [planes[idx] for idx in plane_ids]
	floor_planes = [plane for plane in all_planes if plane['type'] == 'floor']
	assert len(floor_planes) == 1, 'There should be only one floor for each room'
	floor_plane = floor_planes[0]
	floor_plane_id = floor_plane['ID']
	line_ids = np.where(plane_to_line[floor_plane_id])[0].tolist()
	floor_lines_raw = [lines[line_id] for line_id in line_ids]
	floor_lines = list()
	for line_idx, floor_line in enumerate(floor_lines_raw):
	c_id_1, c_id_2 = np.where(line_junc_mat[floor_line['ID']])[0].tolist()
	c1 = tuple(junctions[c_id_1]['coordinate'][:2])
	c2 = tuple(np.array(junctions[c_id_2]['coordinate'][:2]))
	if c1 == c2:
	continue
	floor_lines.append((c1, c2))

	floor_lines = list(set(floor_lines)) # remove duplications
	floor_lines = sort_room_edges(floor_lines)
	room_edges.append(floor_lines)
	return room_edges


	def sort_room_edges(lines):
	cur_id = 0
	picked = [False] * len(lines)
	id_list = [0, ]
	while len(id_list) < len(lines):
	line = lines[cur_id]
	picked[cur_id] = True
	check_ = [(line[1] in other) and not picked[other_idx] for other_idx, other in enumerate(lines)]
	next_ids = np.nonzero(check_)[0]
	try:
	assert len(next_ids) == 1
	except:
	raise WrongRoomError('Invalid room shape')
	next_id = next_ids[0]
	id_list.append(next_id)
	if lines[next_id][1] == line[1]: # swap the two endpoints, to make the loop valid.
	lines[next_id] = (lines[next_id][1], lines[next_id][0])
	cur_id = next_id
	if lines[next_id][1] == lines[0][0]: # already form a closed loop, then skip the remaining lines
	break
	sorted_lines = [lines[idx] for idx in id_list]
	return sorted_lines


	def find_all_colinear_paris(all_room_edges):
	colinear_pairs = list()
	for room_idx, room_edges in enumerate(all_room_edges):
	for edge_idx, edge in enumerate(room_edges):
	for other_room_idx, other_edges in enumerate(all_room_edges):
	if other_room_idx < room_idx:
	continue
	for other_edge_idx, other_edge in enumerate(other_edges):
	if other_room_idx == room_idx and other_edge_idx <= edge_idx:
	continue
	if _check_colinear(edge, other_edge, line_dist_th=8):
	ele1 = (room_idx, edge)
	ele2 = (other_room_idx, other_edge)
	colinear_pairs.append([ele1, ele2])
	return colinear_pairs


	def combine_colinear_edges(colinear_pairs):
	all_colinear_sets = list()
	all_pairs = list(colinear_pairs) # make a copy of the input list
	combined = [False] * len(colinear_pairs)

	while len(all_pairs) > 0:
	colinear_set = _combine_colinear_pairs(0, all_pairs, combined)
	all_colinear_sets.append(colinear_set)
	all_pairs = [all_pairs[i] for i in range(len(all_pairs)) if combined[i] is False]
	combined = [False] * len(all_pairs)
	return all_colinear_sets


	def _combine_colinear_pairs(idx, all_pairs, combined):
	colinear_set = set(all_pairs[idx])
	combined[idx] = True
	for other_idx, pair in enumerate(all_pairs):
	if not combined[other_idx] and (
	all_pairs[idx][0] in all_pairs[other_idx] or all_pairs[idx][1] in all_pairs[other_idx]):
	colinear_set = colinear_set.union(_combine_colinear_pairs(other_idx, all_pairs, combined))
	return colinear_set


	def _check_colinear(e1, e2, line_dist_th=8):
	# first check whether two line segments are parallel to each other, if not, return False directly
	len_e1 = len_edge(e1)
	len_e2 = len_edge(e2)
	# we need to always make e2 the shorter one
	if len_e1 < len_e2:
	e1, e2 = e2, e1
	v1_01 = (e1[1][0] - e1[0][0], e1[1][1] - e1[0][1])
	v1_10 = (e1[0][0] - e1[1][0], e1[0][1] - e1[1][1])
	v2_01 = (e2[1][0] - e2[0][0], e2[1][1] - e2[0][1])
	v2_10 = (e2[0][0] - e2[1][0], e2[0][1] - e2[1][1])
	len_1 = np.sqrt(v1_01[0] 2 + v1_01[1] 2)
	len_2 = np.sqrt(v2_01[0] 2 + v2_01[1] 2)
	if len_1 == 0 or len_2 == 0:
	cos = 0
	else:
	cos = (v1_01[0] * v2_01[0] + v1_01[1] * v2_01[1]) / (len_1 * len_2)
	if abs(cos) > 0.99:
	# then check the distance between two parallel lines
	len_10_20 = len_edge((e1[0], e2[0]))
	len_10_21 = len_edge((e1[0], e2[1]))
	len_11_20 = len_edge((e1[1], e2[0]))
	len_11_21 = len_edge((e1[1], e2[1]))

	# two endpoints are very close, then we can say these two edges are colinear
	if np.min([len_10_20, len_10_21, len_11_20, len_11_21]) <= 5:
	return True
	# otherwise we need to check the distance first
	v_10_20 = (e2[0][0] - e1[0][0], e2[0][1] - e1[0][1])
	cos_11_10_20 = (v1_01[0] * v_10_20[0] + v1_01[1] * v_10_20[1]) / (len_1 * len_10_20)
	sin_11_10_20 = np.sqrt(1 - cos_11_10_20 ** 2)
	dist_20_e1 = len_10_20 * sin_11_10_20
	if dist_20_e1 <= line_dist_th:
	# we need two check whether they have some overlaps
	v_11_20 = (e2[0][0] - e1[1][0], e2[0][1] - e1[1][1])
	cos_10_11_20 = (v1_10[0] * v_11_20[0] + v1_10[1] * v_11_20[1]) / (len_1 * len_11_20)
	if cos_11_10_20 >= 0 and cos_10_11_20 >= 0:
	return True
	v_10_21 = (e2[1][0] - e1[0][0], e2[1][1] - e1[0][1])
	cos_11_10_21 = (v1_01[0] * v_10_21[0] + v1_01[1] * v_10_21[1]) / (len_1 * len_10_21)
	v_11_21 = (e2[1][0] - e1[1][0], e2[1][1] - e1[1][1])
	cos_10_11_21 = (v1_10[0] * v_11_21[0] + v1_10[1] * v_11_21[1]) / (len_1 * len_11_21)
	if cos_11_10_21 >= 0 and cos_10_11_21 >= 0:
	return True
	return False
	else:
	# if the two line segments have distance > 3, we can say they are not colinear
	return False
	else:
	return False


	def merge_edges(edges):
	base_e = edges[0][1]
	merged_edges = [edges[0], ]
	base_len = np.sqrt((base_e[1][0] - base_e[0][0]) 2 + (base_e[1][1] - base_e[0][1]) 2)
	base_unit_v = ((base_e[1][0] - base_e[0][0]) / base_len, (base_e[1][1] - base_e[0][1]) / base_len)

	for edge in edges[1:]:
	room_idx = edge[0]
	e = edge[1]
	v_b0e0 = (e[0][0] - base_e[0][0], e[0][1] - base_e[0][1])
	proj_len = (v_b0e0[0] * base_unit_v[0] + v_b0e0[1] * base_unit_v[1])
	proj_e0 = (int(base_e[0][0] + base_unit_v[0] * proj_len), int(base_e[0][1] + base_unit_v[1] * proj_len))
	proj_e1 = (int(proj_e0[0] + e[1][0] - e[0][0]), int(proj_e0[1] + e[1][1] - e[0][1]))
	new_e = (proj_e0, proj_e1)
	new_edge = (room_idx, new_e)
	merged_edges.append(new_edge)

	adjusted_merged_edges = adjust_colinear_edges(merged_edges)

	return adjusted_merged_edges


	def adjust_colinear_edges(edges):
	base_corner = (edges[0][0], edges[0][1][0])
	all_corners = [base_corner, (edges[0][0], edges[0][1][1])]
	for edge in edges[1:]:
	all_corners.append((edge[0], edge[1][0]))
	all_corners.append((edge[0], edge[1][1]))
	unit_v = unit_v_edge(edges[0][1])
	corner_projs = list()
	# FIXME: need to fix the corner coords here, it's wrong now!! they are not merged...
	for room, other_c in all_corners:
	v_base_c = (other_c[0] - base_corner[1][0], other_c[1] - base_corner[1][1])
	proj = (unit_v[0] * v_base_c[0] + unit_v[1] * v_base_c[1])
	corner_projs.append(proj)
	order = np.argsort(corner_projs).tolist()

	# # merge corners that are too close to the prev corner
	# for o_idx, corner_idx in enumerate(order[1:]):
	# corner = all_corners[corner_idx][1]
	# prev_idx = order[o_idx]
	# prev_corner = all_corners[prev_idx][1]
	# dist = len_edge((corner, prev_corner))
	# if dist <= 5:
	# all_corners[corner_idx] = (all_corners[corner_idx][0], prev_corner)

	adjusted_edges = list()
	for idx, edge in enumerate(edges):
	room_idx = edge[0]
	idx_1 = idx * 2
	idx_2 = idx * 2 + 1
	adj_idx_1 = order.index(idx_1)
	adj_idx_2 = order.index(idx_2)
	step_direction = 1 if adj_idx_2 > adj_idx_1 else -1
	adjusted_edge = list()
	for o_idx in range(adj_idx_1, adj_idx_2, step_direction):
	c_idx = order[o_idx]
	next_c_idx = order[o_idx + step_direction]
	segment = (room_idx, (all_corners[c_idx][1], all_corners[next_c_idx][1]))
	if len_edge(segment[1]) == 0:
	continue
	adjusted_edge.append(segment)
	adjusted_edges.append(adjusted_edge)
	return adjusted_edges


	def adjust_room_edges(room_edges):
	refined_room_edges = list()

	init_room_edges = list(room_edges)
	for edge_i, edge in enumerate(room_edges):
	next_i = edge_i
	while True:
	next_i = next_i + 1 if next_i < len(room_edges) - 1 else 0
	next_edge = room_edges[next_i]
	if next_edge[0] != next_edge[1]:
	break
	if edge[1] == next_edge[0]: # no need for refining
	refined_room_edges.append(edge)
	else: # the two corners disagree, refinement is required
	if edge[0] == edge[1]:
	print('skip collasped edge')
	continue
	unit_edge = unit_v_edge(edge)
	ext_edge = ((edge[0][0] - unit_edge[0] * 50, edge[0][1] - unit_edge[1] * 50),
	(edge[1][0] + unit_edge[0] * 50, edge[1][1] + unit_edge[1] * 50))
	unit_next = unit_v_edge(next_edge)
	ext_next = ((next_edge[0][0] - unit_next[0] * 50, next_edge[0][1] - unit_next[1] * 50),
	(next_edge[1][0] + unit_next[0] * 50, next_edge[1][1] + unit_next[1] * 50))
	intersec = get_intersection(ext_edge[0], ext_edge[1], ext_next[0], ext_next[1])
	try:
	assert intersec is not None
	except:
	print('no intersect, move endpoint directly')
	intersec = next_edge[0]
	intersec = (int(np.round(intersec[0])), int(np.round(intersec[1])))
	refined_edge = (edge[0], intersec)
	refined_room_edges.append(refined_edge)
	room_edges[edge_i] = refined_edge
	room_edges[next_i] = (intersec, next_edge[1])
	if next_i < edge_i:
	refined_room_edges[next_i] = room_edges[next_i]

	# drop collapsed edges
	refined_room_edges = [edge for edge in refined_room_edges if edge[0] != edge[1]]
	for edge_i in range(len(refined_room_edges)):
	next_i = edge_i + 1 if edge_i < len(refined_room_edges) - 1 else 0
	if refined_room_edges[edge_i][1] != refined_room_edges[next_i][0]:
	new_edge = (refined_room_edges[edge_i][0], refined_room_edges[next_i][0])
	refined_room_edges[edge_i] = new_edge
	return refined_room_edges


	def clean_room_edges(all_room_edges):
	refined_room_paths = [_extract_room_path(room_edges) for room_edges in all_room_edges]
	corner_to_room = defaultdict(list)
	for room_idx, room_path in enumerate(refined_room_paths):
	for corner in room_path:
	corner_to_room[corner].append(room_idx)
	# remove I-shape corner used by only one room
	for room_idx, room_edges in enumerate(all_room_edges):
	cp_room_edges = list(room_edges)
	rm_flag = True
	while rm_flag:
	rm_flag = False
	for edge_i, edge in enumerate(cp_room_edges):
	prev_i = edge_i - 1
	prev_edge = cp_room_edges[prev_i]
	if _check_colinear(prev_edge, edge, line_dist_th=5):
	rm_candidate = edge[0]
	if len(corner_to_room[rm_candidate]) == 1 and corner_to_room[rm_candidate][0] == room_idx:
	cp_room_edges[prev_i] = (prev_edge[0], edge[1])
	rm_flag = True
	cp_room_edges.pop(edge_i)
	break
	next_i = edge_i + 1 if edge_i < len(cp_room_edges) - 1 else 0
	next_edge = cp_room_edges[next_i]
	if _check_colinear(next_edge, edge, line_dist_th=5):
	rm_candidate = edge[1]
	if len(corner_to_room[rm_candidate]) == 1 and corner_to_room[rm_candidate][0] == room_idx:
	cp_room_edges[next_i] = (edge[0], next_edge[1])
	rm_flag = True
	cp_room_edges.pop(edge_i)
	break
	if len(cp_room_edges) != len(room_edges):
	all_room_edges[room_idx] = cp_room_edges

	corner_to_room = get_corner_to_room(all_room_edges)
	all_corners = list(corner_to_room.keys())
	corners_to_merge = find_corners_to_merge(all_corners, corner_to_room)
	while corners_to_merge is not None:
	num_aff = [len(corner_to_room[x]) for x in corners_to_merge]
	order = np.argsort(num_aff)[::-1]
	base_corner = corners_to_merge[order[0]]
	for corner in corners_to_merge:
	if corner == base_corner:
	continue
	all_room_edges = move_corner(corner, base_corner, corner_to_room, all_room_edges)

	corner_to_room = get_corner_to_room(all_room_edges)
	all_corners = list(corner_to_room.keys())
	corners_to_merge = find_corners_to_merge(all_corners, corner_to_room)

	# for room_idx, room_edges in enumerate(all_room_edges):
	# cp_room_edges = list(room_edges)
	# rm_flag = True
	# while rm_flag:
	# rm_flag = False
	# for edge_i, edge in enumerate(cp_room_edges):
	# len_e = len_edge(edge)
	# if len_e <= 5:
	# if len(corner_to_room[edge[0]]) == 1:
	# prev_i = edge_i - 1
	# prev_edge = cp_room_edges[prev_i]
	# cp_room_edges[prev_i] = (prev_edge[0], edge[1])
	# rm_flag = True
	# cp_room_edges.pop(edge_i)
	# break
	# elif len(corner_to_room[edge[1]]) == 1:
	# next_i = edge_i + 1 if edge_i < len(cp_room_edges) - 1 else 0
	# next_edge = cp_room_edges[next_i]
	# cp_room_edges[next_i] = (edge[0], next_edge[1])
	# rm_flag = True
	# cp_room_edges.pop(edge_i)
	# else:
	# continue
	#
	# if len(cp_room_edges) != len(room_edges):
	# all_room_edges[room_idx] = cp_room_edges

	return all_room_edges


	def move_corner(c, target, corner_to_room, all_room_edges):
	rooms = corner_to_room[c]
	for room_idx in rooms:
	for edge_idx, edge in enumerate(all_room_edges[room_idx]):
	if c in edge:
	if c == edge[0]:
	new_edge = (target, edge[1])
	elif c == edge[1]:
	new_edge = (edge[0], target)
	else:
	continue
	all_room_edges[room_idx][edge_idx] = new_edge
	return all_room_edges


	def find_corners_to_merge(all_corners, corner_to_room, th=5):
	all_close_pairs = list()
	for idx1, corner in enumerate(all_corners):
	for idx2, other_corner in enumerate(all_corners):
	if idx2 <= idx1:
	continue
	if len_edge((corner, other_corner)) <= th:
	rooms_1 = tuple(sorted(corner_to_room[corner]))
	rooms_2 = tuple(sorted(corner_to_room[other_corner]))
	if rooms_1 == rooms_2:
	continue
	elif len(rooms_1) ==1:
	if rooms_1[0] in list(rooms_2):
	continue
	else:
	all_close_pairs.append([corner, other_corner])
	elif len(rooms_2) ==1:
	if rooms_2[0] in list(rooms_1):
	continue
	else:
	all_close_pairs.append([corner, other_corner])
	else:
	all_close_pairs.append([corner, other_corner])

	if len(all_close_pairs) == 0:
	return None

	close_set = find_one_close_set(all_close_pairs)
	corners_to_merge = list(close_set)

	return corners_to_merge


	def find_one_close_set(all_corner_paris):
	all_pairs = list(all_corner_paris) # make a copy of the input list
	combined = [False] * len(all_corner_paris)

	close_set = _combine_colinear_pairs(0, all_pairs, combined)

	return close_set


	def get_corner_to_room(all_room_edges):
	room_paths = [_extract_room_path(room_edges) for room_edges in all_room_edges]
	corner_to_room = defaultdict(list)
	for room_idx, room_path in enumerate(room_paths):
	for corner in room_path:
	corner_to_room[corner].append(room_idx)
	return corner_to_room


	def filter_rooms(all_room_edges, im_size):
	# filter rooms that are covered by larger rooms
	room_masks = list()
	updated_room_edges = list()
	for room_edges in all_room_edges:
	room_mask = draw_room_seg_from_edges(room_edges, im_size)
	if room_mask is not None and room_mask.sum() > 20: # remove too small rooms
	room_masks.append(room_mask)
	updated_room_edges.append(room_edges)
	all_room_edges = updated_room_edges

	removed = list()
	for room_idx, room_mask in enumerate(room_masks):
	# do not consider the current room, and do not consider removed rooms
	other_masks = [room_masks[i] for i in range(len(all_room_edges)) if i != room_idx and i not in removed]
	if len(other_masks) == 0: # if all other masks are removed..
	other_masks_all = np.zeros([im_size, im_size])
	else:
	other_masks_all = np.clip(np.sum(np.stack(other_masks, axis=-1), axis=-1), 0, 1)
	joint_mask = np.clip(other_masks_all + room_mask, 0, 1)
	mask_area = room_mask.sum()
	overlap_area = mask_area + other_masks_all.sum() - joint_mask.sum()
	if overlap_area / mask_area > 0.5:
	removed.append(room_idx)

	all_room_edges = [all_room_edges[idx] for idx in range(len(all_room_edges)) if idx not in removed]

	return all_room_edges


	## Utils

	class WrongRoomError(Exception):
	pass

	def _extract_room_path(room_edges):
	room_path = [edge[0] for edge in room_edges]
	return room_path


	def len_edge(e):
	return np.sqrt((e[1][0] - e[0][0]) 2 + (e[1][1] - e[0][1]) 2)


	def unit_v_edge(e):
	len_e = len_edge(e)
	assert len_e != 0
	unit_v = ((e[1][0] - e[0][0]) / len_e, (e[1][1] - e[0][1]) / len_e)
	return unit_v


	def get_intersection(p0, p1, p2, p3):
	"""
	reference: StackOverflow https://stackoverflow.com/questions/563198/how-do-you-detect-where-two-line-segments-intersect#565282
	"""
	s1_x = p1[0] - p0[0]
	s1_y = p1[1] - p0[1]
	s2_x = p3[0] - p2[0]
	s2_y = p3[1] - p2[1]

	s = (-s1_y * (p0[0] - p2[0]) + s1_x * (p0[1] - p2[1])) / (-s2_x * s1_y + s1_x * s2_y)
	t = (s2_x * (p0[1] - p2[1]) - s2_y * (p0[0] - p2[0])) / (-s2_x * s1_y + s1_x * s2_y)

	if 1 >= s >= 0 and 1 >= t >= 0:
	i_x = p0[0] + (t * s1_x)
	i_y = p0[1] + (t * s1_y)
	return (i_x, i_y)
	else:
	return None


	def draw_room_seg_from_edges(edges, im_size):
	edge_map = np.zeros([im_size, im_size])
	for edge in edges:
	edge = np.array(edge).astype(np.int)
	cv2.line(edge_map, tuple(edge[0]), tuple(edge[1]), 1, 3)
	reverse_edge_map = 1 - edge_map
	label, num_features = ndimage.label(reverse_edge_map)
	if num_features < 2:
	return None
	bg_label = label[0, 0]
	num_labels = [(label==l).sum() for l in range(1, num_features+1)]
	num_labels[bg_label-1] = 0
	room_label = np.argmax(num_labels) + 1
	room_map = np.zeros([im_size, im_size])
	room_map[np.where(label == room_label)] = 1

	return room_map



	if __name__ == '__main__':
	data_base = './montefloor_data/'
	dir_names = list(sorted(os.listdir(data_base)))

	invalid_scenes = list()

	for dir_name in dir_names:
	if 'scene' not in dir_name:
	continue
	data_dir = os.path.join(data_base, dir_name)
	annot_path = os.path.join(data_dir, 'annotation_3d.json')
	with open(annot_path) as f:
	annot = json.load(f)
	image_path = os.path.join(data_dir, 'density.png')

	try:
	generate_graph(annot, image_path, data_dir)
	except WrongRoomError:
	invalid_scenes.append(dir_name)
	print('Finish processing data {}'.format(dir_name))

	print('Failed on {} scenes with invalid rooms: {}'.format(len(invalid_scenes), invalid_scenes))