1. collect_data.py

Overview

The collect_data.py script is mostly unchanged from the version I forked from Sentdex. It is used to collect the raw training data for all model versions. At this point, any differences in the training data fed into each model version is done by additional preprocessing scripts, so there is only one collect data script at this time. Anything of note about this script, which I don't cover here can probably be learned about in Sentdex's Tutorial for the original version of this script.

Dependencies

This script requires the below imports in order to run. The grabscreen.py and getkeys.py scripts are available in this repo and you'll have to install pypiwin32, numpy and opencv if you don't already have them.

import numpy as np
from grabscreen import grab_screen
import cv2
import time
from getkeys import key_check
import os

Configuration Settings

Before you run the script insure the following configuration variables are set to your liking:

Keys Inputs

This script is currently set to only capture nine possible WASD inputs, but can be reconfigured to capture any of the keyboard inputs noted in the getkeys.py script. Below is an example of how you could modify this script to capture presses of the "P" key. (Note that for each new input you add, you'll have to increase the shape of the multi-hot array by 1 in all areas that it appears.)

w = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
s = [0, 1, 0, 0, 0, 0, 0, 0, 0, 0]
a = [0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
d = [0, 0, 0, 1, 0, 0, 0, 0, 0, 0]
wa = [0, 0, 0, 0, 1, 0, 0, 0, 0, 0]
wd = [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
sa = [0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
sd = [0, 0, 0, 0, 0, 0, 0, 1, 0, 0]
nk = [0, 0, 0, 0, 0, 0, 0, 0, 1, 0]
p = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1]

You'll also have to update the keys_to_output() function to look for the new key in the key list generated by the getkeys script. (Note that how you order the If/Else statements for each key will decide which key saved to the training data file. i.e. If W, A and P are pressed at the same time, the function shown below would record the input as just WA because its if statement comes before P's elif statement.

def keys_to_output(keys):
    '''
    Convert keys to a ...multi-hot... array
     0  1  2  3  4   5   6   7    8     9
    [W, S, A, D, WA, WD, SA, SD, NOKEY, P] boolean values.
    '''
    output = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    if 'W' in keys and 'A' in keys:
        output = wa
    elif 'W' in keys and 'D' in keys:
        output = wd
    elif 'S' in keys and 'A' in keys:
        output = sa
    elif 'S' in keys and 'D' in keys:
        output = sd
    elif 'W' in keys:
        output = w
    elif 'S' in keys:
        output = s
    elif 'A' in keys:
        output = a
    elif 'D' in keys:
        output = d
    elif 'P' in keys:
        output = p
    else:
        output = nk
    return output

Training File Starting Value and Data Cap

The starting_value variable is used in the file name when saving your .npy training files and the data_cap variable is to set how many frames will be collected between each save. Each time a file is saved the script will increment the starting_value up by 1. (Note that in my experience, saving 25,000 480px by 270px RGB images with multi hot labels will create around a 12GB training file, so make sure you only save files to a size that your PC can load into memory.)

starting_value = 1

data_cap = 25000

Note that there is a if/else statement place inside the first while loop. It ensures that existing training files are not overwritten if you forget to update the starting_value variable.

Game Window Properties

The game_width, game_height and title_bar variables are fed into the grab_screen() function and should match the window dimensions of your game in pixels. The script is assuming the game window is positioned in the upper right hand corner of your screen. Most title bars are 40px tall, so you probably won't have to change title_bar very often.

game_width = 1920
game_height = 1080
title_bar = 40

Training Image Dimensions

The width, height variables specify in pixels, the dimensions your training images will be rescaled to and saved as. (Doubling both these values is a great way of quadrupling the size of your training data files.)

width = 480
height = 270

Training Data File Paths

file_name variable should match the path and filename format that you want you want your training data saved to. Please note that this var will need to be updated in two places when changing the path. Not very pythonic and eventually I'll look into fixing it by adding an additional path variable thats fed into the file_name variable.

The first location file_name

while True:
    file_name = 'D:/training_data/raw_data/training_data-{}.npy'.format(starting_value)

    if os.path.isfile(file_name):
        print('File exists, moving along', starting_value)
        starting_value += 1
    else:
        print('File does not exist, starting fresh!', starting_value)

        break

The second location file_name

                if len(training_data) == data_cap:
                    print('SAVING..')
                    np.save(file_name, training_data)
                    print('SAVED')
                    training_data = []
                    starting_value += 1
                    file_name = 'D:/training_data/raw_data/training_data-{}.npy'.format(starting_value)

Running The Script

After you have installed the required dependencies and set the config variables to your liking, run the script as you would any other python script. Running it from Command Prompt is supposably fastest, but I run it from Pycharm and still get enough frames per second to get the job done. Upon launching the script, it will print a short countdown allowing you to ALT + TAB over to your game window before it starts recording data. While it's recording you can press "T" to pause the recording if you need to. Pausing will initiate an auto save, so do not try to unpause until after the console prints "SAVED!"