rrt.py

#!/usr/bin/env python

# rrtstar.py
# This program generates a 
# asymptotically optimal rapidly exploring random tree (RRT* proposed by Sertac Keraman, MIT) in a rectangular region.
#
# Originally written by Steve LaValle, UIUC for simple RRT in
# May 2011
# Modified by Md Mahbubur Rahman, FIU for RRT* in
# January 2016

import sys, random, math, pygame
from pygame.locals import *
from math import sqrt,cos,sin,atan2
from lineIntersect import *
import time 

#constants
XDIM = 640
YDIM = 480
WINSIZE = [XDIM, YDIM]
EPSILON = 7.0
NUMNODES = 2000
RADIUS=15
OBS=[(500,150,100,50),(300,80,100,50),(150,220,100,50)]

def obsDraw(pygame,screen):
    blue=(0,0,255)
    for o in OBS: 
      pygame.draw.rect(screen,blue,o)

def dist(p1,p2):
    return sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1]))

def step_from_to(p1,p2):
    if dist(p1,p2) < EPSILON:
        return p2
    else:
        theta = atan2(p2[1]-p1[1],p2[0]-p1[0])
        return p1[0] + EPSILON*cos(theta), p1[1] + EPSILON*sin(theta)

def chooseParent(nn,newnode,nodes):
        for p in nodes:
         if checkIntersect(p,newnode,OBS) and dist([p.x,p.y],[newnode.x,newnode.y]) <dist([nn.x,nn.y],[newnode.x,newnode.y]):
          nn = p
        newnode.cost=nn.cost+dist([nn.x,nn.y],[newnode.x,newnode.y])
        newnode.parent=nn
        return newnode,nn

def path_length(nodes):
  length = 0
  for i in range(len(nodes)-1):
    length+=dist([nodes[i].x,nodes[i].y],[nodes[i+1].x,nodes[i+1].y])
  return length

def drawSolutionPath(start,goal,nodes,pygame,screen):
	pink = 200, 20, 240
	nn = nodes[0]
	for p in nodes:
	   if dist([p.x,p.y],[goal.x,goal.y]) < dist([nn.x,nn.y],[goal.x,goal.y]):
	      nn = p
	while nn!=start:
		pygame.draw.line(screen,pink,[nn.x,nn.y],[nn.parent.x,nn.parent.y],5)  
		nn=nn.parent


def get_path(start,goal,nodes):
  ret_nodes = []
  nn = nodes[0]
  for p in nodes:
     if dist([p.x,p.y],[goal.x,goal.y]) < dist([nn.x,nn.y],[goal.x,goal.y]):
        nn = p
  while nn!=start:
    ret_nodes.append(nn)
    nn=nn.parent
  return ret_nodes

class Node:
    x = 0
    y = 0
    cost=0  
    parent=None
    def __init__(self,xcoord, ycoord):
         self.x = xcoord
         self.y = ycoord
	
def main(imgno):
    #initialize and prepare screen
    #a=checkIntersect()
    #print(a)
    pygame.init()
    screen = pygame.display.set_mode(WINSIZE)
    pygame.display.set_caption('RRTstar')
    white = 255, 255, 255
    black = 20, 20, 40
    screen.fill(white)
    obsDraw(pygame,screen)


    t = time.time()

    nodes = []
    
    #nodes.append(Node(XDIM/2.0,YDIM/2.0)) # Start in the center
    nodes.append(Node(0.0,0.0)) # Start in the corner
    start=nodes[0]
    goal=Node(630.0,470.0)
    for i in range(NUMNODES):
        rand = Node(random.random()*XDIM, random.random()*YDIM)
        nn = nodes[0]
        for p in nodes:
          if dist([p.x,p.y],[rand.x,rand.y]) < dist([nn.x,nn.y],[rand.x,rand.y]):
            nn = p
        interpolatedNode= step_from_to([nn.x,nn.y],[rand.x,rand.y])
	
        newnode = Node(interpolatedNode[0],interpolatedNode[1])
        if checkIntersect(nn,rand,OBS):
          
          [newnode,nn]=chooseParent(nn,newnode,nodes);
       
          nodes.append(newnode)
          pygame.draw.line(screen,black,[nn.x,nn.y],[newnode.x,newnode.y])
          # nodes=reWire(nodes,newnode,pygame,screen)
          pygame.display.update()
        #print i, "    ", nodes

          for e in pygame.event.get():
            if e.type == QUIT or (e.type == KEYUP and e.key == K_ESCAPE):
                sys.exit("Leaving because you requested it.")
    elapsed = time.time()-t
    expanded = len(nodes)
    path = get_path(start,goal,nodes)
    path_len = path_length(path)
    drawSolutionPath(start,goal,nodes,pygame,screen)
    pygame.display.update()
    print expanded,elapsed,path_len
    pygame.image.save(screen, "rrt_stats/image"+str(imgno)+".jpg")
    pygame.quit()
# if python says run, then we should run
if __name__ == '__main__':
  for i in range(10):
    main(i)
    # running = True
    # while running:
    #    for event in pygame.event.get():
    #        if event.type == pygame.QUIT:
    #              running = False