class Actor(object):
def __init__(self):
# Start session
self.sess = tf.InteractiveSession()
self.sess.run(tf.global_variables_initializer())
# Load Model
saver = tf.train.Saver()
saver.restore(self.sess, "./model.ckpt")
# Init contoller for manual override
self.real_controller = XboxController()
def get_action(self, obs):
### determine manual override
manual_override = self.real_controller.LeftBumper == 1
# -- Automatic Joystick
## Look
vec = resize_image(obs)
## Think
joystick = \
model.y.eval(session=self.sess, feed_dict={model.x: [vec], model.keep_prob: 1.0})[0]
# -- Real Joystick
joystick_real = self.real_controller.read()
joystick_real[1] *= -1 # flip y (this is in the config when it runs normally)
## Act
### calibration
output = [
int(joystick[0] * 80),
int(joystick[1] * 80),
int(round(joystick[2])),
int(round(joystick[3])),
int(round(joystick[4])),
0
]
output_real = [
int(joystick_real[0] * 80),
int(joystick_real[1] * 80),
int(round(joystick_real[2])),
int(round(joystick_real[3])),
int(round(joystick_real[4])),
int(self.real_controller.LeftBumper)
]
### print to console
cprint("Player: " + str(output_real), 'yellow')
cprint("AI: " + str(output), 'green')
return [output, output_real]
def main(file_name, starting_value):
# Instantiate the gamepad
controller = XboxController()
file_name = file_name
starting_value = starting_value
training_data = []
# Countdown
for i in list(range(4))[::-1]:
print(i + 1)
time.sleep(1)
last_time = time.time()
paused = False
print("STARTING RECORDING!!!!!!!!!")
while True:
if not paused:
screen = grab_screen(region=(0, 40, 1920, 1040))
last_time = time.time()
# resize to something a bit more acceptable for a CNN
screen = cv2.resize(screen, (480, 270))
# run a color convert:
screen = cv2.cvtColor(screen, cv2.COLOR_BGR2RGB)
controller_output = controller.read()
training_data.append([screen, controller_output])
last_time = time.time()
if len(training_data) % SAMPLING_RATE == 0:
print(len(training_data))
# Save a checkpoint
if len(training_data) == SAVE_RATE:
np.save(file_name, training_data)
print("SAVED TRAINING DATA")
training_data = []
starting_value += 1
file_name = 'data/training_data-{}.npy'.format(
starting_value)
keys = key_check()
# if 'S' in keys:
# np.save(file_name, training_data)
# print('SAVED TRAINING DATA')
if 'T' in keys:
if paused:
paused = False
print('unpaused!')
time.sleep(1)
else:
print('Pausing!')
paused = True
time.sleep(1)
class Actor(object):
def __init__(self):
# Load in model from train.py and load in the trained weights
self.model = create_model(keep_prob=1) # no dropout
self.model.load_weights('model_weights.h5')
# Init contoller for manual override
self.real_controller = XboxController()
def get_action(self, obs):
### determine manual override
manual_override = self.real_controller.LeftBumper == 1
go=self.real_controller.RightBumper==1
if not manual_override:
## Look
vec = resize_image(obs)
vec = np.expand_dims(vec, axis=0) # expand dimensions for predict, it wants (1,66,200,3) not (66, 200, 3)
## Think
joystick = self.model.predict(vec, batch_size=1)[0]
else:
joystick = self.real_controller.read()
joystick[1] *= -1 # flip y (this is in the config when it runs normally)
## Act
### calibration
if manual_override:
output = [
(int(joystick[0] * 10000)-40)*2,
int(joystick[1] * 10000)-40,
int(joystick[2]),
int(round(joystick[3])),
int(round(joystick[4])),
]
else:
output = [
round(joystick[0])*40,
0,
int(joystick[2]+1),
0,
0,
]
### print to console
if manual_override:
cprint("Manual: " + str(output), 'yellow')
else:
cprint("AI: " + str(output), 'green')
if not go and not manual_override:
output[2]=0
return output
class Actor(object):
def __init__(self):
self.model = create_model(keep_prob=1)
self.model.load_weights(MODEL_NAME)
# Init controller for manual override
self.real_controller = XboxController()
def get_action(self, screen):
manual_override = self.real_controller.UpDPad == 1
if not manual_override:
# Look
vec = screen
vec = np.expand_dims(vec, axis=0)
# Think
ai_control = self.model.predict(vec, batch_size=1)[0]
else:
joystick = self.real_controller.read()
joystick[1] *= -1
### calibration
output = [
round((ai_control[0]), 3), # L / R
round((ai_control[1]), 3), # U / D
round((ai_control[2]), 3),
round((ai_control[3]), 3),
int(round(ai_control[4])), # A
int(round(ai_control[5])), # B
int(round(ai_control[6])), # X
int(round(ai_control[7])), # Y
int(round(ai_control[8])), # B
int(round(ai_control[9])), # T
]
### print to console
if manual_override:
print("Manual: " + str(output))
else:
print("AI: " + str(output))
return output
class Actor(object):
def __init__(self):
# Load in model from train.py and load in the trained weights
self.model = create_model(keep_prob=1) # no dropout
self.model.load_weights('model_weights.h5')
# Init contoller for manual override
self.real_controller = XboxController()
def get_action(self, obs):
### determine manual override
manual_override = self.real_controller.LeftBumper == 1
if not manual_override:
## Look
vec = resize_image(obs)
vec = np.expand_dims(vec, axis=0) # expand dimensions for predict, it wants (1,66,200,3) not (66, 200, 3)
## Think
joystick = self.model.predict(vec, batch_size=1)[0]
else:
joystick = self.real_controller.read()
joystick[1] *= -1 # flip y (this is in the config when it runs normally)
## Act
### calibration
output = [
int(joystick[0] * 80),
int(joystick[1] * 80),
int(round(joystick[2])),
int(round(joystick[3])),
int(round(joystick[4])),
]
pressed = ''
if output[0] > 30:
keyboard.press('right')
pressed += 'right '
if output[0] < -30:
keyboard.press('left')
pressed += 'left '
if output[1] > 30:
keyboard.press('up')
pressed += 'up '
if output[1] < -30:
keyboard.press('down')
pressed += 'down '
if output[2] > 0.5:
keyboard.press('z')
pressed += 'z '
if output[3] > 0.5:
keyboard.press('z')
pressed += 'z '
### print to console
if manual_override:
print("Manual: " + str(output)+ "Pressed: " + pressed)
else:
print("AI: " + str(output)+ "Pressed: " + pressed)
return output
class MainWindow(wx.Frame):
""" Main frame of the application
"""
title = 'Data Acquisition'
def __init__(self):
wx.Frame.__init__(self, None, title=self.title, size=(660, 330))
# Init controller
self.controller = XboxController()
# Create GUI
self.create_main_panel()
# Timer
self.timer = wx.Timer(self)
self.Bind(wx.EVT_TIMER, self.on_timer, self.timer)
self.rate = SAMPLE_RATE
self.idle_rate = IDLE_SAMPLE_RATE
self.timer.Start(self.idle_rate)
self.recording = False
self.t = 0
def create_main_panel(self):
# Panels
self.img_panel = wx.Panel(self)
self.joy_panel = wx.Panel(self)
self.record_panel = wx.Panel(self)
# Images
img = wx.Image(320, 240)
self.image_widget = wx.StaticBitmap(self.img_panel, wx.ID_ANY,
wx.Bitmap(img))
# Joystick
self.init_plot()
self.PlotCanvas = FigCanvas(self.joy_panel, wx.ID_ANY, self.fig)
# Recording
self.txt_outputDir = wx.TextCtrl(self.record_panel,
wx.ID_ANY,
pos=(5, 0),
size=(320, 30))
uid = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
self.txt_outputDir.ChangeValue("samples/" + uid)
self.btn_record = wx.Button(self.record_panel,
wx.ID_ANY,
label="Record",
pos=(335, 0),
size=(100, 30))
self.Bind(wx.EVT_BUTTON, self.on_btn_record, self.btn_record)
self.Bind(wx.EVT_UPDATE_UI, self.on_update_btn_record, self.btn_record)
# sizers
sizer = wx.BoxSizer(wx.HORIZONTAL)
sizer.Add(self.img_panel, 0, wx.ALL, 5)
sizer.Add(self.joy_panel, 0, wx.ALL, 5)
mainSizer_v = wx.BoxSizer(wx.VERTICAL)
mainSizer_v.Add(sizer, 0, wx.ALL, 5)
mainSizer_v.Add(self.record_panel, 0, wx.ALL, 5)
# finalize layout
self.SetAutoLayout(True)
self.SetSizer(mainSizer_v)
self.Layout()
def init_plot(self):
self.plotMem = 50 # how much data to keep on the plot
self.plotData = [[0] * (5)] * self.plotMem # mem storage for plot
self.fig = Figure((4, 3))
self.axes = self.fig.add_subplot(111)
def on_timer(self, event):
self.poll()
# stop drawing if recording to avoid slow downs
#if self.recording == False:
# self.draw()
def poll(self):
self.bmp = take_screenshot()
self.controller_data = self.controller.read()
self.update_plot()
if self.recording == True:
self.save_data()
def update_plot(self):
self.plotData.append(
self.controller_data) # adds to the end of the list
self.plotData.pop(
0) # remove the first item in the list, ie the oldest
def save_data(self):
image_file = self.outputDir + '/' + 'img_' + str(self.t) + '.png'
self.bmp.SaveFile(image_file, wx.BITMAP_TYPE_PNG)
# make / open outfile
outfile = open(self.outputDir + '/' + 'data.csv', 'a')
# write line
outfile.write(image_file + ',' +
','.join(map(str, self.controller_data)) + '\n')
outfile.close()
self.t += 1
def draw(self):
# Image
img = self.bmp.ConvertToImage()
img = img.Rescale(320, 240)
self.image_widget.SetBitmap(img.ConvertToBitmap())
# Joystick
x = np.asarray(self.plotData)
self.axes.plot(range(0, self.plotMem), x[:, 0], 'r')
self.axes.hold(True)
self.axes.plot(range(0, self.plotMem), x[:, 1], 'b')
self.axes.plot(range(0, self.plotMem), x[:, 2], 'g')
self.axes.plot(range(0, self.plotMem), x[:, 3], 'k')
self.axes.plot(range(0, self.plotMem), x[:, 4], 'y')
self.axes.hold(False)
self.PlotCanvas.draw()
def on_update_btn_record(self, event):
label = "Stop" if self.recording else "Record"
self.btn_record.SetLabel(label)
def on_btn_record(self, event):
# pause timer
self.timer.Stop()
# switch state
self.recording = not self.recording
if self.recording:
self.start_recording()
# un pause timer
if self.recording:
self.timer.Start(self.rate)
else:
self.timer.Start(self.idle_rate)
def start_recording(self):
# check that a dir has been specified
if self.txt_outputDir.IsEmpty():
msg = wx.MessageDialog(self, 'Specify the Output Directory',
'Error', wx.OK | wx.ICON_ERROR)
msg.ShowModal() == wx.ID_YES
msg.Destroy()
self.recording = False
else: # a directory was specified
self.outputDir = self.txt_outputDir.GetValue()
self.t = 0
# check if path exists - ie may be saving over data
if os.path.exists(self.outputDir):
msg = wx.MessageDialog(
self, 'Output Directory Exists - Overwrite Data?',
'Yes or No', wx.YES_NO | wx.ICON_QUESTION)
result = msg.ShowModal() == wx.ID_YES
msg.Destroy()
# overwrite the data
if result == True:
# delete the dir
shutil.rmtree(self.outputDir)
# re-make dir
os.mkdir(self.outputDir)
# do not overwrite the data
else: # result == False
self.recording = False
self.txt_outputDir.SetFocus()
# no directory so make one
else:
os.mkdir(self.outputDir)
def on_exit(self, event):
self.Destroy()
class MainWindow():
""" Main frame of the application
"""
def __init__(self):
self.root = tk.Tk()
self.sct = mss.mss()
self.root.title('Data Acquisition')
self.root.geometry("660x325")
self.root.resizable(False, False)
# Init controller
self.controller = XboxController()
# Create GUI
self.create_main_panel()
# Timer
self.rate = IDLE_SAMPLE_RATE
self.sample_rate = SAMPLE_RATE
self.idle_rate = IDLE_SAMPLE_RATE
self.recording = False
self.t = 0
self.pause_timer = False
self.on_timer()
self.root.mainloop()
def create_main_panel(self):
# Panels
top_half = tk.Frame(self.root)
top_half.pack(side=tk.TOP, expand=True, padx=5, pady=5)
message = tk.Label(
self.root, text="(Note: UI updates are disabled while recording)")
message.pack(side=tk.TOP, padx=5)
bottom_half = tk.Frame(self.root)
bottom_half.pack(side=tk.LEFT, padx=5, pady=10)
# Images
self.img_panel = tk.Label(top_half,
image=ImageTk.PhotoImage(
"RGB", size=IMAGE_SIZE)) # Placeholder
self.img_panel.pack(side=tk.LEFT, expand=False, padx=5)
# Joystick
self.init_plot()
self.PlotCanvas = FigCanvas(figure=self.fig, master=top_half)
self.PlotCanvas.get_tk_widget().pack(side=tk.RIGHT,
expand=False,
padx=5)
# Recording
textframe = tk.Frame(bottom_half, width=332, height=15, padx=5)
textframe.pack(side=tk.LEFT)
textframe.pack_propagate(0)
self.outputDirStrVar = tk.StringVar()
self.txt_outputDir = tk.Entry(textframe,
textvariable=self.outputDirStrVar,
width=100)
self.txt_outputDir.pack(side=tk.LEFT)
self.outputDirStrVar.set("samples/" +
datetime.now().strftime('%Y-%m-%d_%H:%M:%S'))
self.record_button = ttk.Button(bottom_half,
text="Record",
command=self.on_btn_record)
self.record_button.pack(side=tk.LEFT, padx=5)
def init_plot(self):
self.plotMem = 50 # how much data to keep on the plot
self.plotData = [[0] * (5)] * self.plotMem # mem storage for plot
self.fig = Figure(figsize=(4, 3), dpi=80) # 320,240
self.axes = self.fig.add_subplot(111)
def on_timer(self):
self.poll()
# stop drawing if recording to avoid slow downs
if self.recording == False:
self.draw()
if not self.pause_timer:
self.root.after(self.rate, self.on_timer)
def poll(self):
self.img = self.take_screenshot()
self.controller_data = self.controller.read()
self.update_plot()
if self.recording == True:
self.save_data()
self.t += 1
def take_screenshot(self):
# Get raw pixels from the screen
sct_img = self.sct.grab({
"top": Screenshot.OFFSET_Y,
"left": Screenshot.OFFSET_X,
"width": Screenshot.SRC_W,
"height": Screenshot.SRC_H
})
# Create the Image
return Image.frombytes('RGB', sct_img.size, sct_img.bgra, 'raw',
'BGRX')
def update_plot(self):
self.plotData.append(
self.controller_data) # adds to the end of the list
self.plotData.pop(
0) # remove the first item in the list, ie the oldest
def save_data(self):
image_file = self.outputDir + '/' + 'img_' + str(self.t) + '.png'
self.img.save(image_file)
# write csv line
self.outfile.write(image_file + ',' +
','.join(map(str, self.controller_data)) + '\n')
def draw(self):
# Image
self.img.thumbnail(IMAGE_SIZE, Image.ANTIALIAS) # Resize
self.img_panel.img = ImageTk.PhotoImage(self.img)
self.img_panel['image'] = self.img_panel.img
# Joystick
x = np.asarray(self.plotData)
self.axes.clear()
self.axes.plot(range(0, self.plotMem), x[:, 0], 'r')
self.axes.plot(range(0, self.plotMem), x[:, 1], 'b')
self.axes.plot(range(0, self.plotMem), x[:, 2], 'g')
self.axes.plot(range(0, self.plotMem), x[:, 3], 'k')
self.axes.plot(range(0, self.plotMem), x[:, 4], 'y')
self.PlotCanvas.draw()
def on_btn_record(self):
# pause timer
self.pause_timer = True
if self.recording:
self.recording = False
else:
self.start_recording()
if self.recording:
self.t = 0 # Reset our counter for the new recording
self.record_button["text"] = "Stop"
self.rate = self.sample_rate
# make / open outfile
self.outfile = open(self.outputDir + '/' + 'data.csv', 'a')
else:
self.record_button["text"] = "Record"
self.rate = self.idle_rate
self.outfile.close()
# un pause timer
self.pause_timer = False
self.on_timer()
def start_recording(self):
should_record = True
# check that a dir has been specified
if not self.outputDirStrVar.get():
tkMessageBox.showerror(title='Error',
message='Specify the Output Directory',
parent=self.root)
should_record = False
else: # a directory was specified
self.outputDir = self.outputDirStrVar.get()
# check if path exists - i.e. may be saving over data
if os.path.exists(self.outputDir):
# overwrite the data, yes/no?
if tkMessageBox.askyesno(
title='Warning!',
message='Output Directory Exists - Overwrite Data?',
parent=self.root):
# delete & re-make the dir:
shutil.rmtree(self.outputDir)
os.mkdir(self.outputDir)
# answer was 'no', so do not overwrite the data
else:
should_record = False
self.txt_outputDir.focus_set()
# directory doesn't exist, so make one
else:
os.mkdir(self.outputDir)
self.recording = should_record