Sanmesh's updated collect.py with buttons

collect.py

import os
import math
import numpy as np
import glob
import scipy
import scipy.misc
import datetime
import time as timemod
import RPi.GPIO as GPIO
import argparse

def collectTeleop(num_frames, debug, gpio_collectOff):
    import serial
    from PIL import Image, ImageDraw
    import pygame
    import pygame.camera
    from pygame.locals import *
    pygame.init()
    pygame.camera.init()

    # initialize webcam
    print('initialize webcam')
    cams = pygame.camera.list_cameras()
    cam = pygame.camera.Camera(cams[0],(64,64),'RGB')
    cam.start()

    time_format = '%Y-%m-%d_%H:%M:%S'
    date = datetime.datetime.now()

    imgs_file = 'imgs_{0}'.format(date.strftime(time_format))
    speedx_file = 'speedx_{0}'.format(date.strftime(time_format))
    targets_file = 'targets_{0}'.format(date.strftime(time_format))

    print('connect to serial port')
    if not debug:
        ser = serial.Serial('/dev/ttyACM0')
        if(ser.isOpen() == False):
            ser.open()
    else:
        ser = open('/home/ubuntu/proj/autonomous/test_collect.csv')

    # initialize speeds
    speeds = np.zeros(3,dtype=np.float32)

    # setup storage
    imgs = np.zeros((num_frames,3,64,64),dtype=np.uint8)
    speedx = np.zeros((num_frames,2),dtype=np.float16)
    targets = np.zeros((num_frames,2),dtype=np.float16)

    # Drive in remote control mode
    idx = 0
    GPIO.setmode(GPIO.BCM)
    #gpio_collectOff = 16;
    GPIO.setup(gpio_collectOff, GPIO.IN, pull_up_down=GPIO.PUD_UP) #have separate button for collect off
    input_state = GPIO.input(gpio_collectOff)
    try:
        while input_state != False:
            # Maybe send 'Ok, ready' to the serial port
            # take web footage (every second or whatever)
            img = pygame.surfarray.array3d(cam.get_image())
            ## throw away non-square sides (left and rightmost 20 cols)
            img = img[20:-20]
            ## Shrink to 64x64
            img = scipy.misc.imresize(img,(64,64),'cubic','RGB').transpose(2,1,0)
            # Receive steering + gas (inc. direction)
            if not debug:
                ## Read acceleration information (and time, TODO)
                d = ser.readline()
                ## most recent line
                #data = list(map(int,str(d,'ascii').split(',')))
                line = d.strip()
                data = line.split(b',')
                data = list(map(float,str(d,'ascii').split(',')))
            else:
                d = ser.readline()
                line = d.strip()
                data = list(map(int,line.split(',')))
            # Record all in convenient format
            accel = np.array([float(data[0]),float(data[1]),float(data[2])], dtype=np.float16)
            gx = float(data[3])
            gy = float(data[4])
            gz = float(data[5])
            time = int(data[6])
            steer_raw = int(data[7])
            gas_raw = int(data[8])
            #direction = data[1]
            #if direction:
            #    gas = data[2]
            #else:
            #    gas = -1*data[2]
            #time = data[3]
            #speed = data[4]
            #accel = data[5]
            accel[2] -= 1 # subtract accel due to gravity, maybe the car can fly :p
            # rescale inputs ( decide on max speed and accel of vehicle), clamp values to these
            accel = accel / 10.
            # update speeds, accel is scaled to be m/s**2 at this point
            # so just multiply by seconds elapsed
            #speeds = speeds + accel*(t-t_old)/1000.
            # now shift the accel
            accel += 0.5
            # compute magnitude of speed and accel
            #mspeed = np.sqrt(np.sum(speeds*speeds))
            mspeed = 99
            maccel = np.sqrt(np.sum(accel*accel))

            # Stuff into appropriate arrays
            imgs[idx] = np.array(img,dtype=np.uint8)
            speedx[idx] = np.array([mspeed,maccel],dtype=np.float16)
            targets[idx] = np.array([steer_raw,gas_raw],dtype=np.float16)
            # increment counter, maybe save
            idx += 1
            if idx % num_frames == 0:
                np.savez('collectedData/' + imgs_file,imgs)
                np.savez('collectedData/' + speedx_file,speedx)
                np.savez('collectedData/' + targets_file,targets)
                ## Make new data files
                date = datetime.datetime.now()
                imgs_file = 'imgs_{0}'.format(date.strftime(time_format))
                speedx_file = 'speedx_{0}'.format(date.strftime(time_format))
                targets_file = 'targets_{0}'.format(date.strftime(time_format))
                # zero out storage
                imgs[:] = 0
                speedx[:] = 0
                targets[:] = 0
            timemod.sleep(1)
            input_state = GPIO.input(gpio_collectOff) #check state of button to see if exit
    except:
        np.savez('collectedData/' + imgs_file,imgs[:idx])
        np.savez('collectedData/' + speedx_file,speedx[:idx])
        np.savez('collectedData/' + targets_file,targets[:idx])
        raise      
    return
"""
parser = argparse.ArgumentParser(description='Steer Otto, the autonomous tractor and collect data.')
parser.add_argument('-n','--num_frames', action='store', default=100)
parser.add_argument('-d','--debug', action='store_true', default=False)

args = parser.parse_args()
num_frames = int(args.num_frames)
debug = args.debug
"""