def predict(image, model, device, threshold=None, output_file=None, show_image=False):
transformed_image = get_transform()(image)
model_output = model([transformed_image.to(device)])
prediction = {
"boxes": model_output[0]["boxes"].tolist(),
"confidence_scores": model_output[0]["scores"].tolist(),
"labels": [
REVERSED_CLASS_LABEL_MAP[label]
for label in model_output[0]["labels"].tolist()
],
}
if threshold:
prediction = threshold_scores(prediction, threshold)
if show_image or output_file:
plot_prediction(prediction, image, output_file)
return prediction
def create_mock_image(as_tensor=False):
image = Image.new("RGB", size=(50, 50), color=(155, 0, 0))
if as_tensor:
transform = get_transform()
return transform(image) # convert this to tensor
return image
def create_mock_dataset():
return LISADataset("mock_path", transforms=get_transform())
def training_pipeline(dataset_path):
start_training = perf_counter()
# train on the GPU or on the CPU if a GPU is not available
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
print("Training on the", device)
# Load the dataset. Remove any unwanted images
full_dataset = LISADataset(dataset_path,
transforms=get_transform(),
bad_images=bad_images)
# split data into training, testing and validation sets
data_loader_train, data_loader_test, data_loader_val = split_data(
full_dataset)
validation_ground_truth = create_bounding_boxes(data_loader_val,
BBType.GroundTruth,
model=None,
device=device)
num_classes = len(CLASS_LABEL_MAP) + 1 # add 1 for background class
# hyperparameter tuning
best_parameters, values, experiment, ax_model = optimize(
parameters=[
{
"name": "lr",
"type": "range",
"bounds": [1e-6, 0.01],
"log_scale": True
},
{
"name": "num_epochs",
"type": "choice",
"values": list(range(1, 3))
},
],
evaluation_function=lambda params: train(
params,
num_classes,
device,
data_loader_train,
data_loader_val,
validation_ground_truth,
),
objective_name="mean_average_precision",
total_trials=1,
)
end_training = perf_counter()
training_time = end_training - start_training
print("Training took", training_time / 60.0, "minutes")
# evaluation
start_eval = perf_counter()
path_to_best_model = "model_lr_{}_epochs_{}.pth".format(
best_parameters["lr"], best_parameters["num_epochs"])
print("Best model saved at", path_to_best_model)
model = torch.load(path_to_best_model, map_location=torch.device(device))
testing_ground_truth = create_bounding_boxes(data_loader_test,
BBType.GroundTruth, model,
device)
print("getting evaluation scores...")
evaluation = evaluate(model, data_loader_test, device,
testing_ground_truth)
# TODO: move map calculation into evaluation, just return m_a_p
precisions = [
0 if np.isnan(metric["AP"]) else metric["AP"] for metric in evaluation
]
mean_average_precision = np.sum(precisions) / len(CLASS_LABEL_MAP)
print("Mean average precision:", mean_average_precision)
end_eval = perf_counter()
eval_time = end_eval - start_eval
print("Evaluation took", eval_time / 60.0, "minutes")
def test_removes_bad_images(self):
smaller_dataset = LISADataset("mock_path",
transforms=get_transform(),
bad_images=["nightClip2--0001.jpg"])
self.assertEqual(smaller_dataset.__len__(), 2)
def setUpClass(self):
self.dataset = LISADataset("mock_path", transforms=get_transform())