Main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec  9 09:06:43 2020

@author: jelmer
"""

import math
import numpy as np
from PIL import Image
import random
import os.path
import pygame as pg

from create_object_borders import create_object_borders
from sample_points import sample_points
from nearest_neighb import near_neighb
from edge_check import edges_in_free_space
from Dijkstra import Dijkstra
from splines import bspline_path
from PurePursuit_PID import control
import simulation_plot_functions as spf

pg.init()
running = True
cos, sin = np.cos, np.sin
random.seed(70)
cm_to_pix = 100/70                                                              #Conversion rate

#%% Import floor_plan, create borders around objects and save pixels in array:
Map = "A"                                                                       #Change to B or C for other map
border_size = 25                                                                #Define border size in pixels
image_name = "floor_plan_" + Map + "_" + str(border_size) + ".png"
array_name = "floor_plan_" + Map + "_" + str(border_size) + ".npy"
    
if os.path.isfile(image_name) and os.path.isfile(array_name):                   #Load if files exist
    floor_plan = Image.open(image_name)
    pixels = np.load(array_name)
else:     
    file = "floor_plan_" + Map + ".png"                                         #Create border image and array if not exist
    floor_plan = Image.open(file)
    pixels, floor_plan = create_object_borders(floor_plan, border_size)
    floor_plan.convert('RGB').save(image_name)
    np.save(array_name, pixels)
    print("Borders Created")

width, height = floor_plan.size                                                 #Define floor_plan shape

#%% Sample random points:
D = 30                                                                          #Devide factor (larger is less sampled points)

coords, N_S = sample_points(width, height, pixels, D)                           #Store coordinates
print("Number of sampled points is: ", N_S)
        
#%% Search neigbours within treshold:
DIST = 400*(math.log(N_S)/N_S)**(1/3)                                           #Treshold distance
NN = near_neighb(coords[:,0], coords[:,1], DIST)                                #Store NN for all points
print("NN found")                     

#%% Check which edges between neighbours are in free space:
connections = edges_in_free_space(NN, coords, pixels)                           #Store connections in free space
print("Connections found")

#%% Find shortest path with A*:
start = 100
end = 200

route = Dijkstra(coords, connections, start, end)                               #Find path with A*

#%% Smoothening of the path:
n_edges = len(route)-1                                                          #number of edges in the route
sn = 30*n_edges                                                                 #Number of samples to be plotted in the spline path
k=5                                                                             #bspline polynomial degree
cx, cy, dx, dy, ddx, ddy = bspline_path(coords[route][:,0], coords[route][:,1], sn, k)    #bspline path

#%% Pure pursuit and PID controller:
init_pose = np.array([0,0, np.arctan2(dy[0],dx[0])])                            #Initial pose of robot
tx, ty, yaw, delta, v, F = control(cx, cy, init_pose)                           #Robot path generated by controller

#%% Simulation:

#Create screen with floor_plan:
screen = pg.display.set_mode((width,height))
floor_plan = pg.image.load(image_name)

#Initials of robot:
base_l = 30 * cm_to_pix
base_w = 30 * cm_to_pix
base_init = pg.Surface((base_l, base_w), pg.SRCALPHA)
base_init.fill((0, 0, 255, 100))

wheel_l = 15 * cm_to_pix
wheel_w = 10 * cm_to_pix
front_wheel_init = pg.Surface((wheel_l, wheel_w), pg.SRCALPHA)
front_wheel_init.fill((0, 0, 0, 100))

rear_wheel_init = pg.Surface((wheel_l, wheel_w), pg.SRCALPHA)
rear_wheel_init.fill((0, 0, 0, 100))

#Settings and initials:
t = 0
frame = 0
screenshot = 0
record = False
m_loc = None
ptx, pty = None, None

while running:
    #Show background image:
    screen.blit(floor_plan, (0, 0))
    
    #Plot nodes in yellow:
    spf.plot_nodes(coords, screen)
    
    #Plot connections in cyan:
    spf.plot_connections(connections, coords, screen)
    
    #Plot original route in blue:
    spf.plot_original_route(route, coords, screen)

    #Plot smoothed route in red:
    spf.plot_smoothed_route(cx, cy, screen)
            
    #Plot travelled route:
    spf.plot_travelled_route(t, tx, ty, screen)
    
    #Plot privious travelled route:
    if False:
        if ptx is not None:
            spf.plot_privious_travelled_route(ptx, pty, screen)
        
    #Define positions and angles:
    yaw_rad = yaw[t]
    yaw_deg = np.rad2deg(yaw_rad)
    steer_rad = delta[t]
    steer_deg = np.rad2deg(steer_rad)

    center = np.array([tx[t], ty[t]])
    rear = center - base_l/2 * np.array([cos(yaw_rad), sin(yaw_rad)])
    front = center + base_l/2 * np.array([cos(yaw_rad), sin(yaw_rad)])
    
    #Plotting front wheel:
    spf.plot_front_wheel(front_wheel_init, yaw_deg, steer_deg, front, screen)
    
    #Plotting rear wheel:
    spf.plot_rear_wheel(rear_wheel_init, yaw_deg, rear, screen)
    
    #Plotting base:
    spf.plot_base(base_init, yaw_deg, center, screen)
        
    #Go to next time-step:
    if t < min(len(tx),len(delta))-1:
        t += 1
        
    #Mouse:
    if any(pg.mouse.get_pressed()) == True :
        m_loc = pg.mouse.get_pos()
    
    #Draw next location:    
    if m_loc != None:
        node = np.argmin(np.linalg.norm(coords-np.asarray(m_loc), axis=1))
        pg.draw.circle(screen, (255, 0, 0), tuple(coords[node]), 4)
        
    
    #Determine new path:
    if t >= min(len(tx),len(delta))-1 and m_loc != None:
        
        try:
            #Copy travelled route to privious travelled route:
            ptx, pty = tx[:t], ty[:t]
            
            #Set new start and end and determine new route:
            start = end
            end = node
            route = Dijkstra(coords, connections, start, end)
            
            #Set current position as begin:
            route_x = np.append(center[0], coords[route][:,0])
            route_y = np.append(center[1], coords[route][:,1])
            
            #Determnine new smooth path and P-controller path:
            cx, cy, dx, dy, ddx, ddy = bspline_path(route_x, route_y, sn, k)
            init_pose = np.array([0,0, yaw_rad])                                   
            tx, ty, yaw, delta, v, F = control(cx, cy, init_pose)
            
            #Reset t:
            t = 0
        
        except:
            print("Error occurred, please try again!")                        
        
        #Reset location:
        m_loc = None
        
    #Save simulation:
    if record == True:
        filename = "frames/" + str(frame) + ".png"
        pg.image.save(screen, filename)
        frame += 1
        
    #Update screen:
    pg.display.update()
    pg.time.delay(75)
    
    #Events:
    for event in pg.event.get():
        
        #Close screen:
        if event.type==pg.QUIT:
            running = False
            
        #Make screenshot:
        if event.type == pg.KEYDOWN:
            if event.key == pg.K_s:
                filename = "screenshots/" + str(screenshot) + ".png"
                pg.image.save(screen, filename)
                screenshot +=1
 
pg.quit()