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