tracepath.py

#!/usr/bin/env python

import Image
import sys

class Node(object):
	'''This class is a 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())