def get_polygons(nodes, rgb_im):
polygons = []
unvisited = [node for node in nodes if node.visited is False]
while unvisited:
poly = polygon.Polygon()
current = unvisited[0]
current.visited = True
closest, distance = get_closest_unvisited_node(current, nodes, rgb_im)
if distance is not None and distance > constants.MAX_NODE_DIFFERENCE:
unvisited = unvisited[1:]
continue
while closest is not None:
poly.nodes.append(current)
current = closest
current.visited = True
closest, distance = get_closest_unvisited_node(current, nodes, rgb_im)
if closest is None:
break
i = -1
while distance > constants.MAX_NODE_DIFFERENCE:
if (current is poly.nodes[0] or i < -constants.MAX_NODE_BACKTRACK
or (utils.get_distance(poly.nodes[0], current) < constants.MAX_NODE_DIFFERENCE)):
closest = None
break
current = poly.nodes[i]
closest, distance = get_closest_unvisited_node(current, unvisited, rgb_im)
i -= 1
if len(poly.nodes) > 2:
polygons.append(poly)
unvisited = [node for node in nodes if node.visited is False]
return polygons
def get_closest_unvisited_node(current, nodes, rgb_im):
closest_node = None
shortest_distance = None
pos = nodes.index(current)
i = 1
go_up = True
go_down = True
while (0 <= pos - i or len(nodes) > pos + i) and (go_up or go_down):
for sign in [-1, 1]:
if sign == -1 and not go_down or sign == 1 and not go_up:
continue
index = pos + i*sign
if not 0 <= index < len(nodes):
continue
node = nodes[index]
if closest_node is not None:
if sign == -1 and shortest_distance < current.x - node.x:
go_down = False
elif sign == 1 and shortest_distance < node.x - current.x:
go_up = False
if node.visited:
continue
distance = utils.get_distance(nodes[pos], node)
distance *= utils.get_water_multiplier(current, node, rgb_im)
if shortest_distance is None or distance < shortest_distance:
closest_node = node
shortest_distance = distance
i += 1
return closest_node, shortest_distance
def prune_overlapping_nodes(polygons):
assert polygons
polygons = utils.sort_by_polygon_length(polygons)
polygons.reverse()
exterior_polygons = [polygons[0]]
for test_polygon in polygons[1:]:
starting_count = len(test_polygon.nodes)
for exterior_polygon in exterior_polygons:
exterior_nodes = test_polygon.get_exterior_nodes(exterior_polygon)
if not exterior_nodes:
if len(test_polygon.nodes) == starting_count:
exterior_polygon.inner = test_polygon
elif (exterior_polygon is exterior_polygons[-1] and
len(exterior_nodes) > 2 and
utils.get_distance(exterior_nodes[0], exterior_nodes[-1]) <=
constants.MAX_NODE_DIFFERENCE):
test_polygon.nodes = exterior_nodes
exterior_polygons.append(test_polygon)
break
return exterior_polygons
def prune_extra_nodes(polygons):
pruned_polygons = []
for poly in polygons:
if len(poly.nodes) < constants.MINIMUM_PRUNING_SIZE:
assert len(poly.nodes) > 2
pruned_polygons.append(poly)
continue
pruned = polygon.Polygon(poly.nodes[:2])
for node in poly.nodes[2:]:
if utils.get_distance(pruned.nodes[-1], node) <= 1:
continue
end_node = None
if utils.same_line(pruned.nodes[-2], pruned.nodes[-1], node):
end_node = node
else:
if end_node is None:
end_node = node
pruned.nodes.append(end_node)
if len(pruned.nodes) > 2:
pruned_polygons.append(pruned)
return pruned_polygons
