def main():
parser = argparse.ArgumentParser(
description="Steering prediction with PilotNet")
parser.add_argument("--config_file",
type=str,
required=True,
help="Config file to use")
args = parser.parse_args()
config_fp = args.config_file
config = load_config(config_fp)
#mandatory parameters
learning_rate = float(config['learning_rate'])
root_dir, annotation_file = os.path.split(config['annotation_file'])
prefix, _ = annotation_file.split(".")
#optional parameters
file_prefix = config['file_prefix']
checkpoint_file = config['checkpoint_file']
load_files = bool(config['load_files'])
use_float32 = bool(config['use_float32'])
optical_flow = bool(config['optical_flow'])
label_scale = float(config['label_scale'])
momentum = float(config['momentum'])
batch_size = int(config['batch_size'])
gpu = int(config['gpu'])
epochs = int(config['epochs'])
workers = int(config['workers'])
context_length = int(config['context_length'])
sequence_length = int(config['sequence_length'])
hidden_dim = int(config['hidden_dim'])
_, config_file = os.path.split(config_fp)
config_file_name, _ = config_file.split(".")
output_dir = config_file_name.replace("\n", "")
prefix = prefix + file_prefix + 'commandant'
network = models.CommandantNet(context_length=context_length,
sequence_length=sequence_length,
hidden_dim=hidden_dim,
use_float32=use_float32,
gpu=gpu,
optical_flow=optical_flow)
im_size = (125, 400)
starting_epoch = 0
if (checkpoint_file != ''):
dir, file = os.path.split(checkpoint_file)
_, ep = file.split("epoch")
num, ext = ep.split(".")
starting_epoch = int(num)
print("Starting Epoch number:", starting_epoch)
state_dict = torch.load(checkpoint_file)
network.load_state_dict(state_dict)
if (label_scale == 1.0):
label_transformation = None
else:
label_transformation = transforms.Compose(
[transforms.Lambda(lambda inputs: inputs.mul(label_scale))])
if (gpu >= 0):
network = network.cuda(gpu)
if (use_float32):
network.float()
trainset = loaders.F1SequenceDataset(
root_dir,
annotation_file,
im_size,
context_length=context_length,
sequence_length=sequence_length,
use_float32=True,
label_transformation=label_transformation,
optical_flow=optical_flow)
else:
network.double()
trainset = loaders.F1SequenceDataset(
root_dir,
annotation_file,
im_size,
context_length=context_length,
sequence_length=sequence_length,
use_float32=False,
label_transformation=label_transformation,
optical_flow=optical_flow)
# trainset.read_files()
if optical_flow:
if (load_files or
(not os.path.isfile("./" + prefix + "_commandantopticalflows.pkl"))
or
(not os.path.isfile("./" + prefix +
"_commandantopticalflowannotations.pkl"))):
trainset.read_files_flow()
trainset.write_pickles(
prefix + "_commandantopticalflows.pkl",
prefix + "_commandantopticalflowannotations.pkl")
else:
trainset.read_pickles(
prefix + "_commandantopticalflows.pkl",
prefix + "_commandantopticalflowannotations.pkl")
else:
if (load_files or (not os.path.isfile("./" + prefix + "_images.pkl"))
or (not os.path.isfile("./" + prefix + "_annotations.pkl"))):
trainset.read_files()
trainset.write_pickles(prefix + "_images.pkl",
prefix + "_annotations.pkl")
else:
trainset.read_pickles(prefix + "_images.pkl",
prefix + "_annotations.pkl")
''' '''
mean, stdev = trainset.statistics()
mean_ = torch.from_numpy(mean)
stdev_ = torch.from_numpy(stdev)
if use_float32:
mean_.float()
stdev_.float()
trainset.img_transformation = transforms.Normalize(mean_, stdev_)
config['image_transformation'] = trainset.img_transformation
config['label_transformation'] = trainset.label_transformation
print("Using configuration: ", config)
if (not os.path.isdir(output_dir)):
os.makedirs(output_dir)
config_dump = open(os.path.join(output_dir, "config.pkl"), 'wb')
pickle.dump(config, config_dump)
config_dump.close()
trainLoader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=workers)
print(trainLoader)
#Definition of our loss.
criterion = nn.MSELoss()
# Definition of optimization strategy.
optimizer = optim.SGD(network.parameters(),
lr=learning_rate,
momentum=momentum)
train_model(network,
criterion,
optimizer,
trainLoader,
prefix,
output_dir,
n_epochs=epochs,
gpu=gpu,
starting_epoch=starting_epoch)
def main():
parser = argparse.ArgumentParser(description="Test AdmiralNet")
parser.add_argument("--model_file", type=str, required=True)
parser.add_argument("--annotation_file", type=str, required=True)
parser.add_argument("--write_images", action="store_true")
parser.add_argument("--plot", action="store_true")
args = parser.parse_args()
plot = args.plot
annotation_dir, annotation_file = os.path.split(args.annotation_file)
model_dir, model_file = os.path.split(args.model_file)
config_path = os.path.join(model_dir, 'config.pkl')
config_file = open(config_path, 'rb')
config = pickle.load(config_file)
print(config)
model_prefix, _ = model_file.split(".")
# return
gpu = int(config['gpu'])
optical_flow = bool(config.get('optical_flow', ''))
use_float32 = bool(config['use_float32'])
label_scale = float(config['label_scale'])
prefix, _ = annotation_file.split(".")
prefix = prefix + config['file_prefix'] + 'commandant'
context_length = int(config['context_length'])
sequence_length = int(config['sequence_length'])
hidden_dim = int(config['hidden_dim'])
size = (125, 400)
network = models.CommandantNet(context_length=context_length,
sequence_length=sequence_length,
hidden_dim=hidden_dim,
use_float32=use_float32,
gpu=gpu,
optical_flow=optical_flow)
state_dict = torch.load(args.model_file)
network.load_state_dict(state_dict)
print(network)
if (label_scale == 1.0):
label_transformation = None
else:
label_transformation = transforms.Compose(
[transforms.Lambda(lambda inputs: inputs.mul(label_scale))])
if (use_float32):
network.float()
trainset = loaders.F1SequenceDataset(
annotation_dir,
annotation_file,
size,
context_length=context_length,
sequence_length=sequence_length,
use_float32=True,
label_transformation=label_transformation,
optical_flow=optical_flow)
else:
network.double()
trainset = loaders.F1SequenceDataset(
annotation_dir,
annotation_file,
size,
context_length=context_length,
sequence_length=sequence_length,
label_transformation=label_transformation,
optical_flow=optical_flow)
if (gpu >= 0):
network = network.cuda(gpu)
if optical_flow:
if ((not os.path.isfile("./" + prefix + "_commandantopticalflows.pkl"))
or
(not os.path.isfile("./" + prefix +
"_commandantopticalflowannotations.pkl"))):
trainset.read_files_flow()
trainset.write_pickles(
prefix + "_commandantopticalflows.pkl",
prefix + "_commandantopticalflowannotations.pkl")
else:
trainset.read_pickles(
prefix + "_commandantopticalflows.pkl",
prefix + "_commandantopticalflowannotations.pkl")
else:
if ((not os.path.isfile("./" + prefix + "_images.pkl"))
or (not os.path.isfile("./" + prefix + "_annotations.pkl"))):
trainset.read_files()
trainset.write_pickles(prefix + "_images.pkl",
prefix + "_annotations.pkl")
else:
trainset.read_pickles(prefix + "_images.pkl",
prefix + "_annotations.pkl")
''' '''
mean, stdev = trainset.statistics()
mean_ = torch.from_numpy(mean)
stdev_ = torch.from_numpy(stdev)
if use_float32:
mean_.float()
stdev_.float()
trainset.img_transformation = config['image_transformation']
if plot:
batch_size = 1
else:
batch_size = 32
loader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=False,
num_workers=0)
cum_diff = 0.0
t = tqdm(enumerate(loader))
network.eval()
predictions = []
ground_truths = []
losses = []
criterion = nn.MSELoss()
if (gpu >= 0):
criterion = criterion.cuda(gpu)
if args.write_images:
imdir = "admiralnet_prediction_images_" + model_prefix
os.mkdir(imdir)
annotation_file = open(args.annotation_file, 'r')
annotations = annotation_file.readlines()
annotation_file.close()
im, _, _, _, _ = annotations[0].split(",")
background = cv2.imread(os.path.join(annotation_dir, 'raw_images', im),
cv2.IMREAD_UNCHANGED)
out_size = background.shape
fourcc = cv2.VideoWriter_fourcc(*'XVID')
videoout = cv2.VideoWriter(os.path.join(imdir, "video.avi"), fourcc,
60.0, (out_size[1], out_size[0]), True)
wheel = cv2.imread('steering_wheel.png', cv2.IMREAD_UNCHANGED)
wheelrows = 150
wheelcols = 150
wheel = cv2.resize(wheel, (wheelcols, wheelrows),
interpolation=cv2.INTER_CUBIC)
for idx, (inputs, previous_control, labels) in t:
if (gpu >= 0):
previous_control = previous_control.cuda(gpu)
inputs = inputs.cuda(gpu)
labels = labels.cuda(gpu)
pred = torch.div(network(inputs, previous_control), label_scale)
if plot:
if pred.shape[1] == 1:
angle = pred.item()
ground_truth = labels.item()
else:
angle = pred.squeeze()[0].item()
ground_truth = labels.squeeze()[0].item()
predictions.append(angle)
ground_truths.append(ground_truth)
t.set_postfix(angle=angle, ground_truth=ground_truth)
loss = criterion(pred, labels)
losses.append(torch.mean(loss).item())
# print("Ground Truth: %f. Prediction: %f.\n" %(scaled_ground_truth, scaled_angle))
if args.write_images:
scaled_angle = 180.0 * angle
M = cv2.getRotationMatrix2D((wheelrows / 2, wheelcols / 2),
scaled_angle, 1)
wheel_rotated = cv2.warpAffine(wheel, M, (wheelrows, wheelcols))
numpy_im = np.transpose(trainset.images[idx],
(1, 2, 0)).astype(np.float32)
# print(numpy_im.shape)
im, _, _, _, _ = annotations[idx].split(",")
background = cv2.imread(
os.path.join(annotation_dir, 'raw_images', im),
cv2.IMREAD_UNCHANGED)
out_size = background.shape
#print(background.shape)
overlayed = imutils.annotation_utils.overlay_image(
background, wheel_rotated, int((out_size[1] - wheelcols) / 2),
int((out_size[0] - wheelcols) / 2))
name = "ouput_image_" + str(idx) + ".png"
output_path = os.path.join(imdir, name)
cv2.imwrite(output_path, overlayed)
videoout.write(overlayed)
'''
'''
predictions_array = np.array(predictions)
ground_truths_array = np.array(ground_truths)
diffs = np.subtract(predictions_array, ground_truths_array)
rms = np.sqrt(np.mean(np.array(losses)))
print("RMS Error: ", rms)
if args.plot:
from scipy import stats
import matplotlib.pyplot as plt
t = np.linspace(0, len(loader) - 1, len(loader))
plt.plot(t, predictions_array, 'r')
plt.plot(t, ground_truths_array, 'b')
#plt.plot(t,diffs)
plt.show()