Created
November 16, 2008 17:42
-
-
Save NaPs/25516 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
import Image | |
import sys | |
class Node(object): | |
'''This class is an node of a graph''' | |
def __init__(self, name): | |
self.name = name | |
self.reset() | |
def reset(self): | |
self.previous = None | |
self.traveled = float('inf') | |
def __repr__(self): | |
return '<Node: %s (%s)>' % (self.name, self.traveled) | |
class Edge(object): | |
'''This class is an edge of a graph''' | |
def __init__(self, node1, node2, weight): | |
self.node1 = node1 | |
self.node2 = node2 | |
self.weight = weight | |
def give_other_node(self, node): | |
'''Return the other node''' | |
if node is self.node1: | |
return self.node2 | |
elif node is self.node2: | |
return self.node1 | |
def __repr__(self): | |
return '<Edge: %s -%s- %s>' % (self.node1.name, self.weight, self.node2.name) | |
class Graph(object): | |
'''This class is an entire graph''' | |
def __init__(self): | |
self.nodes = [] | |
self.edges = [] | |
self.edge_start = None | |
def add_node(self, node): | |
if not node in self.nodes: | |
self.nodes.append(node) | |
def add_edge(self, edge): | |
if not edge in self.edges: | |
self.edges.append(edge) | |
def get_or_create_edge(self, node1, node2, weight=1): | |
for edge in self.edges: | |
if (edge.node1 == node1 and edge.node2 == node2) or \ | |
(edge.node1 == node2 and edge.node2 == node1): | |
return edge | |
else: | |
new_edge = Edge(node1, node2, weight) | |
g.add_edge(new_edge) | |
return new_edge | |
def get_or_create_node_by_name(self, node_name): | |
for node in self.nodes: | |
if node_name == node.name: | |
return node | |
else: | |
new_node = Node(node_name) | |
self.add_node(new_node) | |
return new_node | |
def give_connected_edges(self, node): | |
edges = [] | |
for edge in self.edges: | |
if node is edge.node1 or node is edge.node2: | |
edges.append(edge) | |
return edges | |
def init_dijkstra(self, node_start): | |
for node in self.nodes: | |
node.reset() | |
self.node_start = node_start | |
self.node_start.traveled = 0 | |
nodes_not_traveled = list(self.nodes) | |
while nodes_not_traveled: | |
n1 = min(nodes_not_traveled, key=lambda x: x.traveled) | |
nodes_not_traveled.remove(n1) | |
for edge in self.give_connected_edges(n1): | |
n2 = edge.give_other_node(n1) | |
if n2.traveled > n1.traveled + edge.weight: | |
n2.traveled = n1.traveled + edge.weight | |
n2.previous = n1 | |
def best_path_dijkstra(self, edge_end): | |
if self.node_start not in self.nodes: | |
raise ValueError('Dijkstra not initialized') | |
path = [] | |
current_node = edge_end | |
while current_node is not self.edge_start: | |
path.insert(0, current_node) | |
current_node = current_node.previous | |
return path | |
def export_graphviz(self): | |
lines = ['graph {'] | |
for node in self.nodes: | |
lines.append(' "%s";' % node.name) | |
for edge in self.edges: | |
lines.append(' "%s" -- "%s" [label="%s"];' % (edge.node1.name, edge.node2.name, edge.weight)) | |
lines.append('}') | |
return '\n'.join(lines) | |
def get_valid_neighbors(image, px_pos_x, px_pos_y): | |
neighbors_pos_list = [] | |
if px_pos_x - 1 > 0: | |
neighbors_pos_list.append((px_pos_x - 1, px_pos_y)) | |
if px_pos_y - 1 > 0: | |
neighbors_pos_list.append((px_pos_x, px_pos_y - 1)) | |
if px_pos_x + 1 < image.size[0]: | |
neighbors_pos_list.append((px_pos_x + 1, px_pos_y)) | |
if px_pos_y + 1 < image.size[1]: | |
neighbors_pos_list.append((px_pos_x, px_pos_y + 1)) | |
px_map = image.load() | |
neighbors_pos_list_filtered = list(neighbors_pos_list) | |
return [(n_x, n_y) for n_x, n_y in neighbors_pos_list if px_map[n_x, n_y] != (0, 0, 0)] | |
if __name__ == '__main__': | |
# Graph initialization : | |
g = Graph() | |
# Image initialization : | |
original_image = Image.open(sys.argv[1]) | |
resized_image = original_image.resize([x/5 for x in original_image.size]) | |
px_map = resized_image.load() | |
# Image to graph conversion : | |
starting_node = None | |
ending_node = None | |
for x in range(resized_image.size[0]): | |
for y in range(resized_image.size[1]): | |
if px_map[x,y] == (0, 0, 0): # black pixel, ignoring; | |
continue | |
else: | |
current_node = g.get_or_create_node_by_name('%s:%s' % (x, y)) | |
if px_map[x,y] == (0, 0, 255): # Starting px | |
starting_node = current_node | |
elif px_map[x,y] == (255, 0, 0): # Ending px | |
ending_node = current_node | |
for neighbor_x, neighbor_y in get_valid_neighbors(resized_image, x, y): | |
neighbor_node = g.get_or_create_node_by_name('%s:%s' % (neighbor_x, neighbor_y)) | |
g.get_or_create_edge(current_node, neighbor_node) | |
# Best path : | |
g.init_dijkstra(starting_node) | |
path = g.best_path_dijkstra(ending_node) | |
for node in path[1:-1]: | |
# Getting px pos : | |
x, y = [int(x) for x in node.name.split(':')] | |
px_map[x, y] = (0, 255, 0) | |
# Resizing image and saving... | |
resized_image = resized_image.resize([x*5 for x in resized_image.size]) | |
resized_image.save(sys.argv[1] + '_solved.png') | |
fg = open(sys.argv[1] + '.dot', 'w') | |
fg.write(g.export_graphviz()) | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment