def main():
parser = argparse.ArgumentParser(description="Steering prediction with PilotNet")
parser.add_argument("--config_file", type=str, required=True, help="Config file to use")
parser.add_argument("--ros", action="store_true")
args = parser.parse_args()
config_fp = args.config_file
config = load_config(config_fp)
#mandatory parameters
learning_rate = float(config['learning_rate'])
annotation_dir = config['annotation_directory']
dataset_type = config['dataset_type']
#optional parameters
file_prefix = config['file_prefix']
load_files = (config['load_files'] == 'True')
load_pickles = (config['load_pickles'] == 'True')
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'])
output_dimension = int(config['output_dimension'])
hidden_dimension = int(config['hidden_dimension'])
config_file = os.path.basename(config_fp)
print(config_file)
config_file_name, _ = config_file.split(".")
output_dir = config_file_name.replace("\n","") + "_" + dataset_type
if(not os.path.isdir(output_dir)):
os.mkdir(output_dir)
size=(66,200)
if args.ros:
files= {'ros_train_blur.csv', 'ros_train_raw.csv', 'ros_train_flip.csv', 'ros_train_bright.csv', 'ros_train_dark.csv'}
else:
files= {'fullview_linear_raw.csv', 'fullview_linear_brightenned.csv', 'fullview_linear_darkenned.csv', 'fullview_linear_flipped.csv'}
# 'fullview_linear_darkenned.csv',\
# 'fullview_linear_darkflipped.csv',\
# 'fullview_linear_brightenned.csv')
datasets = []
for f in files:
absolute_filepath = os.path.join(annotation_dir,f)
if dataset_type=='optical_flow':
input_channels=2
ds = loaders.F1OpticalFlowDataset(absolute_filepath, size, context_length = context_length, sequence_length = sequence_length)
elif dataset_type=='raw_images':
input_channels=1
ds = loaders.F1ImageSequenceDataset(absolute_filepath, size, context_length = context_length, sequence_length = sequence_length)
elif dataset_type=='combined':
input_channels=3
ds = loaders.F1CombinedDataset(absolute_filepath, size, context_length = context_length, sequence_length = sequence_length)
else:
raise NotImplementedError('Dataset of type: ' + dataset_type + ' not implemented.')
datasets.append( ds )
network = models.AdmiralNet_V2(gpu=gpu,context_length = context_length, sequence_length = sequence_length,\
hidden_dim=hidden_dimension, output_dimension = output_dimension, input_channels=input_channels)
if(gpu>=0):
network = network.cuda(gpu)
print(network)
for dataset in datasets:
splits = dataset.annotation_filename.split(".")
prefix = splits[0]
image_pickle = os.path.join(dataset.root_folder,prefix + dataset.image_pickle_postfix)
labels_pickle = os.path.join(dataset.root_folder,prefix + dataset.label_pickle_postfix)
if(os.path.isfile(image_pickle) and os.path.isfile(labels_pickle)):
print("Loading pickles for: " + dataset.annotation_filename)
dataset.loadPickles()
else:
print("Loading files for: " + dataset.annotation_filename)
dataset.loadFiles()
dataset.writePickles()
# if(gpu>=0 and dataset_type=='raw_images'):
# dataset = dataset.cuda(gpu)
trainset = torch.utils.data.ConcatDataset(datasets)
'''
mean,stdev = trainset.statistics()
print(mean)
print(stdev)
img_transformation = transforms.Compose([transforms.Normalize(mean,stdev)])
trainset.img_transformation = img_transformation
'''
trainLoader = torch.utils.data.DataLoader(trainset, batch_size = batch_size, shuffle = True, num_workers = 0)
print(trainLoader)
#Definition of our loss.
criterion = nn.MSELoss()
# Definition of optimization strategy.
optimizer = optim.SGD(network.parameters(), lr = learning_rate, momentum=momentum)
config['image_transformation'] = None
config['input_channels'] = input_channels
print("Using config:")
print(config)
config_dump = open(os.path.join(output_dir,"config.pkl"), 'w+b')
pickle.dump(config,config_dump)
config_dump.close()
# torch.save( self.images, open( os.path.join( self.root_folder, imgname ), 'w+b' ) )
'''
torch.save( network.projector_input, open(os.path.join(output_dir,"projector_input.pt"), 'w+b') )
torch.save( network.init_hidden, open(os.path.join(output_dir,"init_hidden.pt"), 'w+b') )
torch.save( network.init_cell, open(os.path.join(output_dir,"init_cell.pt"), 'w+b') )
'''
train_model(network, criterion, optimizer, trainLoader, output_dir, output_dimension, n_epochs = epochs, gpu = gpu)
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("--gpu", type=int, default=0)
parser.add_argument("--write_images", action="store_true")
parser.add_argument("--plot", action="store_true")
args = parser.parse_args()
annotation_dir, annotation_file = os.path.split(args.annotation_file)
model_dir, model_file = os.path.split(args.model_file)
leaf_model_dir = os.path.basename(model_dir)
config_path = os.path.join(model_dir, 'config.pkl')
config_file = open(config_path, 'rb')
config = pickle.load(config_file)
print(config)
model_prefix = leaf_model_dir + model_file.split(".")[0]
# return
gpu = args.gpu
use_float32 = bool(config['use_float32'])
label_scale = float(config['label_scale'])
size = (66, 200)
prefix, _ = annotation_file.split(".")
prefix = prefix + config['file_prefix']
context_length = int(config['context_length'])
sequence_length = int(config['sequence_length'])
hidden_dim = int(config['hidden_dimension'])
dataset_type = config['dataset_type']
output_dimension = 1
if dataset_type == 'optical_flow':
input_channels = 2
valset = loaders.F1OpticalFlowDataset(args.annotation_file,
size,
context_length=context_length,
sequence_length=sequence_length)
elif dataset_type == 'raw_images':
input_channels = 1
valset = loaders.F1ImageSequenceDataset(
args.annotation_file,
size,
context_length=context_length,
sequence_length=sequence_length)
elif dataset_type == 'combined':
input_channels = 3
valset = loaders.F1CombinedDataset(args.annotation_file,
size,
context_length=context_length,
sequence_length=sequence_length)
#optical_flow = bool(config.get('optical_flow',''))
rnn_cell_type = 'lstm'
network = models.AdmiralNet_V2(input_channels=input_channels,
gpu=gpu,
cell=rnn_cell_type,
context_length=context_length,
sequence_length=sequence_length,
hidden_dim=hidden_dim)
state_dict = torch.load(args.model_file)
network.load_state_dict(state_dict)
#network.projector_input = torch.load( open(os.path.join(model_dir,"projector_input.pt"), 'r+b') ).cuda(gpu)
#network.init_hidden = torch.load( open(os.path.join(model_dir,"init_hidden.pt"), 'r+b') ).cuda(gpu)
# network.init_cell = torch.load( open(os.path.join(model_dir,"init_cell.pt"), 'r+b') ).cuda(gpu)
network = network.float()
print(network)
if (gpu >= 0):
network = network.cuda(gpu)
annotation_prefix = annotation_file.split(".")[0]
image_pickle = os.path.join(valset.root_folder,
prefix + valset.image_pickle_postfix)
labels_pickle = os.path.join(valset.root_folder,
prefix + valset.label_pickle_postfix)
if (os.path.isfile(image_pickle) and os.path.isfile(labels_pickle)):
valset.loadPickles()
else:
valset.loadFiles()
valset.writePickles()
predictions = []
ground_truths = []
losses = []
criterion = nn.MSELoss()
cum_loss = 0.0
num_samples = 0
if (gpu >= 0):
criterion = criterion.cuda(gpu)
network.eval()
batch_size = 1
loader = torch.utils.data.DataLoader(valset,
batch_size=batch_size,
shuffle=False,
num_workers=0)
t = tqdm(enumerate(loader))
for (i, (inputs, labels)) in t:
labels = labels[:, :, 0:output_dimension]
if gpu >= 0:
inputs = inputs.cuda(gpu)
labels = labels.cuda(gpu)
# Forward pass:
outputs = network(inputs)
loss = criterion(outputs, labels)
prediction = ((outputs.cpu())[0][0][0]).item()
ground_truth = ((labels.cpu())[0][0][0]).item()
predictions.append(prediction)
ground_truths.append(ground_truth)
# logging information
loss_ = loss.item()
cum_loss += loss_
num_samples += batch_size
t.set_postfix(cum_loss=cum_loss / num_samples)
predictions_array = np.array(predictions)
ground_truths_array = np.array(ground_truths)
output_file_prefix = "admiralnetv2_prediction_" + os.path.basename(
annotation_dir) + model_prefix
log_name = output_file_prefix + ".txt"
# imdir = "admiralnet_prediction_images_" + model_prefix
#if(os.path.exists(imdir)==False):
# os.mkdir(imdir)
log_output_path = os.path.join(model_dir, log_name)
log = list(zip(ground_truths_array, predictions_array))
print("Writing results to: " + log_output_path)
with open(log_output_path, "w") as myfile:
for x in log:
log_item = [x[0], x[1]]
myfile.write("{0},{1}\n".format(log_item[0], log_item[1]))
diffs = np.square(np.subtract(predictions_array, ground_truths_array))
rms = np.sqrt(np.mean(np.array(diffs)))
rms = np.sqrt(np.mean(np.array(diffs)))
rms_scikit = sqrt(
mean_squared_error(predictions_array, ground_truths_array))
print("RMS Error: ", rms)
print("RMS Error scikit: ", rms_scikit)
if args.plot:
fig = plt.figure()
ax = plt.subplot(111)
t = np.linspace(0, len(predictions_array) - 1, len(predictions_array))
ax.plot(t, predictions_array, 'r', label='Predicted')
ax.plot(t, ground_truths_array, 'b', label='Ground Truth')
ax.legend()
plt.savefig(output_file_prefix + ".jpeg")