Skip to content

Instantly share code, notes, and snippets.

@tikurahul
Created Mar 7, 2020
Embed
What would you like to do?
Multi-tub training
#!/usr/bin/env python3
'''
Scripts to train a keras model using tensorflow.
Basic usage should feel familiar: python train_v2.py --model models/mypilot
Usage:
train.py [--tubs=tub1,tub2] (--model=<model>) [--type=(linear|inferred|tensorrt_linear|tflite_linear)]
Options:
-h --help Show this screen.
'''
import os
import random
from pathlib import Path
import cv2
import numpy as np
from docopt import docopt
from PIL import Image
from sklearn.model_selection import train_test_split
from tensorflow.python.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.python.keras.utils.data_utils import Sequence
import donkeycar
from donkeycar.parts.keras import KerasInferred
from donkeycar.parts.tub_v2 import Tub
from donkeycar.utils import get_model_by_type, load_scaled_image_arr
class TubDataset(object):
'''
Loads the dataset, and creates a train/test split.
'''
def __init__(self, tub_paths, test_size=0.2, shuffle=True):
self.tub_paths = tub_paths
self.test_size = test_size
self.shuffle = shuffle
self.tubs = [Tub(tub_path) in self.tub_paths]
self.records = list()
def train_test_split(self):
print('Loading tub from paths %s' % (self.tub_paths))
for tub in self.tubs:
self.records.extend(tub)
return train_test_split(self.records, test_size=self.test_size, shuffle=self.shuffle)
class TubSequence(Sequence):
def __init__(self, keras_model, images_base_path, config, records=list()):
self.keras_model = keras_model
self.images_base_path = images_base_path
self.config = config
self.records = records
self.batch_size = self.config.BATCH_SIZE
def __len__(self):
return len(self.records) // self.batch_size
def __getitem__(self, index):
count = 0
records = []
images = []
angles = []
throttles = []
is_inferred = type(self.keras_model) is KerasInferred
while count < self.batch_size:
i = (index * self.batch_size) + count
if i >= len(self.records):
break
record = self.records[i]
record = self._transform_record(record)
records.append(record)
count += 1
for record in records:
image = record['cam/image_array']
angle = record['user/angle']
throttle = record['user/throttle']
images.append(image)
angles.append(angle)
throttles.append(throttle)
X = np.array(images)
if is_inferred:
Y = np.array(angles)
else:
Y = [np.array(angles), np.array(throttles)]
return X, Y
def _transform_record(self, record):
for key, value in record.items():
if key == 'cam/image_array' and isinstance(value, str):
image_path = os.path.join(self.images_base_path, value)
image = load_scaled_image_arr(image_path, self.config)
record[key] = image
return record
class ImagePreprocessing(Sequence):
'''
A Sequence which wraps another Sequence with an Image Augumentation.
'''
def __init__(self, sequence, augmentation):
self.sequence = sequence
self.augumentation = augmentation
def __len__(self):
return len(self.sequence)
def __getitem__(self, index):
X, Y = self.sequence[index]
return self.augumentation.augment_images(X), Y
@classmethod
def region_of_interest(cls, image, x, y, width, height):
'''
Selects a region of interest from the original image.
'''
copy = image[y: y + height, x : x + width]
return copy
@classmethod
def crop(cls, left, right, top, bottom, keep_size=False):
'''
The image augumentation sequence.
Crops based on a region of interest among other things.
'''
import imgaug as ia
import imgaug.augmenters as iaa
augmentation = iaa.Sequential([
iaa.Crop(
px=(left, right, top, bottom),
keep_size=keep_size
),
])
return augmentation
@classmethod
def trapezoidal_mask(cls, lower_left, lower_right, upper_left, upper_right, min_y, max_y):
'''
Uses a binary mask to generate a trapezoidal region of interest.
Especially useful in filtering out uninteresting features from an
input image.
'''
import imgaug as ia
import imgaug.augmenters as iaa
def _transform_images(images, random_state, parents, hooks):
# Transform a batch of images
transformed = []
mask = None
for image in images:
if mask is None:
mask = np.zeros(image.shape, dtype='bool')
# # # # # # # # # # # # #
# ul ur min_y
#
#
#
# ll lr max_y
points = [
[upper_left, min_y],
[upper_right, min_y],
[lower_right, max_y],
[lower_left, max_y]
]
cv2.fillConvexPoly(mask, np.array(points, dtype=np.int32), 1)
masked = cv2.bitwise_and(image, mask)
transformed.append(masked)
return transformed
def _transform_keypoints(keypoints_on_images, random_state, parents, hooks):
# No-op
return keypoints_on_images
augmentation = iaa.Sequential([
iaa.Lambda(func_images=_transform_images, func_keypoints=_transform_keypoints)
])
return augmentation
def train(cfg, tub_paths, output_path, model_type):
'''
Train the model
'''
if 'linear' in model_type:
train_type = 'linear'
else:
train_type = model_type
kl = get_model_by_type(train_type, cfg)
kl.compile()
if cfg.PRINT_MODEL_SUMMARY:
print(kl.model.summary())
batch_size = cfg.BATCH_SIZE
dataset = TubDataset(tub_paths, test_size=(1. - cfg.TRAIN_TEST_SPLIT))
images_base_path = dataset.tub.images_base_path
training_records, validation_records = dataset.train_test_split()
print('Records # Training %s' % (len(training_records)))
print('Records # Validation %s' % (len(validation_records)))
training = TubSequence(kl, images_base_path, cfg, training_records)
validation = TubSequence(kl, images_base_path, cfg, validation_records)
# Setup early stoppage callbacks
callbacks = [
EarlyStopping(monitor='val_loss', patience=cfg.EARLY_STOP_PATIENCE),
ModelCheckpoint(
filepath=output_path,
monitor='val_loss',
save_best_only=True,
verbose=1,
)
]
kl.model.fit_generator(
generator=training,
steps_per_epoch=len(training),
callbacks=callbacks,
validation_data=validation,
validation_steps=len(validation),
epochs=cfg.MAX_EPOCHS,
verbose=cfg.VERBOSE_TRAIN,
workers=1,
use_multiprocessing=False
)
def main():
args = docopt(__doc__)
cfg = donkeycar.load_config()
tubs = args['--tubs']
model = args['--model']
model_type = args['--type']
if not model_type:
model_type = cfg.DEFAULT_MODEL_TYPE
tubs = tubs.split(',')
tub_paths = list()
for tub in tubs:
tub_paths.append(Path(os.path.expanduser(tub)).absolute().as_posix())
output_path = os.path.expanduser(model)
train(cfg, tub_paths, output_path, model_type)
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment