// The contents of this file are in the public domain. See LICENSE_FOR_EXAMPLE_PROGRAMS.txt
/*
This example shows how to train a instance segmentation net using the PASCAL VOC2012
dataset. For an introduction to what segmentation is, see the accompanying header file
dnn_instance_segmentation_ex.h.
Instructions how to run the example:
1. Download the PASCAL VOC2012 data, and untar it somewhere.
http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
2. Build the dnn_instance_segmentation_train_ex example program.
3. Run:
./dnn_instance_segmentation_train_ex /path/to/VOC2012
4. Wait while the network is being trained.
5. Build the dnn_instance_segmentation_ex example program.
6. Run:
./dnn_instance_segmentation_ex /path/to/VOC2012-or-other-images
It would be a good idea to become familiar with dlib's DNN tooling before reading this
example. So you should read dnn_introduction_ex.cpp, dnn_introduction2_ex.cpp,
and dnn_semantic_segmentation_train_ex.cpp before reading this example program.
*/
#include "dnn_instance_segmentation_ex.h"
#include "pascal_voc_2012.h"
#include <iostream>
#include <dlib/data_io.h>
#include <dlib/image_transforms.h>
#include <dlib/dir_nav.h>
#include <iterator>
#include <thread>
using namespace std;
using namespace dlib;
// ----------------------------------------------------------------------------------------
// A single training sample for detection. A mini-batch comprises many of these.
struct det_training_sample
{
matrix<rgb_pixel> input_image;
std::vector<dlib::mmod_rect> mmod_rects;
};
// A single training sample for segmentation. A mini-batch comprises many of these.
struct seg_training_sample
{
matrix<rgb_pixel> input_image;
matrix<float> label_image; // The ground-truth label of each pixel. (+1 or -1)
};
// ----------------------------------------------------------------------------------------
bool is_instance_pixel(const dlib::rgb_pixel& rgb_label)
{
if (rgb_label == dlib::rgb_pixel(0, 0, 0))
return false; // Background
if (rgb_label == dlib::rgb_pixel(224, 224, 192))
return false; // The cream-colored `void' label is used in border regions and to mask difficult objects
return true;
}
// Provide hash function for dlib::rgb_pixel
namespace std {
template <>
struct hash<dlib::rgb_pixel>
{
std::size_t operator()(const dlib::rgb_pixel& p) const
{
return (static_cast<uint32_t>(p.red) << 16)
| (static_cast<uint32_t>(p.green) << 8)
| (static_cast<uint32_t>(p.blue));
}
};
}
struct truth_instance
{
dlib::rgb_pixel rgb_label;
dlib::mmod_rect mmod_rect;
};
std::vector<truth_instance> rgb_label_images_to_truth_instances(
const dlib::matrix<dlib::rgb_pixel>& instance_label_image,
const dlib::matrix<dlib::rgb_pixel>& class_label_image
)
{
std::unordered_map<dlib::rgb_pixel, mmod_rect> result_map;
DLIB_CASSERT(instance_label_image.nr() == class_label_image.nr());
DLIB_CASSERT(instance_label_image.nc() == class_label_image.nc());
const auto nr = instance_label_image.nr();
const auto nc = instance_label_image.nc();
for (int r = 0; r < nr; ++r)
{
for (int c = 0; c < nc; ++c)
{
const auto rgb_instance_label = instance_label_image(r, c);
if (!is_instance_pixel(rgb_instance_label))
continue;
const auto rgb_class_label = class_label_image(r, c);
const Voc2012class& voc2012_class = find_voc2012_class(rgb_class_label);
const auto i = result_map.find(rgb_instance_label);
if (i == result_map.end())
{
// Encountered a new instance
result_map[rgb_instance_label] = rectangle(c, r, c, r);
result_map[rgb_instance_label].label = voc2012_class.classlabel;
}
else
{
// Not the first occurrence - update the rect
auto& rect = i->second.rect;
if (c < rect.left())
rect.set_left(c);
else if (c > rect.right())
rect.set_right(c);
if (r > rect.bottom())
rect.set_bottom(r);
DLIB_CASSERT(i->second.label == voc2012_class.classlabel);
}
}
}
std::vector<truth_instance> flat_result;
flat_result.reserve(result_map.size());
for (const auto& i : result_map) {
flat_result.push_back(truth_instance{
i.first, i.second
});
}
return flat_result;
}
// ----------------------------------------------------------------------------------------
struct truth_image
{
image_info info;
std::vector<truth_instance> truth_instances;
};
std::vector<mmod_rect> extract_mmod_rects(
const std::vector<truth_instance>& truth_instances
)
{
std::vector<mmod_rect> mmod_rects(truth_instances.size());
std::transform(
truth_instances.begin(),
truth_instances.end(),
mmod_rects.begin(),
[](const truth_instance& truth) { return truth.mmod_rect; }
);
return mmod_rects;
}
std::vector<std::vector<mmod_rect>> extract_mmod_rect_vectors(
const std::vector<truth_image>& truth_images
)
{
std::vector<std::vector<mmod_rect>> mmod_rects(truth_images.size());
const auto extract_mmod_rects_from_truth_image = [](const truth_image& truth_image)
{
return extract_mmod_rects(truth_image.truth_instances);
};
std::transform(
truth_images.begin(),
truth_images.end(),
mmod_rects.begin(),
extract_mmod_rects_from_truth_image
);
return mmod_rects;
}
det_bnet_type train_detection_network(
const std::vector<truth_image>& truth_images,
unsigned int det_minibatch_size
)
{
const double initial_learning_rate = 0.1;
const double weight_decay = 0.0001;
const double momentum = 0.9;
const double min_detector_window_overlap_iou = 0.65;
const int target_size = 70;
const int min_target_size = 30;
mmod_options options(
extract_mmod_rect_vectors(truth_images),
target_size, min_target_size,
min_detector_window_overlap_iou
);
options.overlaps_ignore = test_box_overlap(0.5, 0.9);
det_bnet_type det_net(options);
det_net.subnet().layer_details().set_num_filters(options.detector_windows.size());
dlib::pipe<det_training_sample> data(200);
auto f = [&data, &truth_images, target_size, min_target_size](time_t seed)
{
dlib::rand rnd(time(0) + seed);
matrix<rgb_pixel> input_image;
random_cropper cropper;
cropper.set_seed(time(0));
cropper.set_chip_dims(350, 350);
// Usually you want to give the cropper whatever min sizes you passed to the
// mmod_options constructor, or very slightly smaller sizes, which is what we do here.
cropper.set_min_object_size(target_size - 2, min_target_size - 2);
cropper.set_max_rotation_degrees(2);
det_training_sample temp;
while (data.is_enabled())
{
// Pick a random input image.
const auto random_index = rnd.get_random_32bit_number() % truth_images.size();
const auto& truth_image = truth_images[random_index];
// Load the input image.
load_image(input_image, truth_image.info.image_filename);
// Get a random crop of the input.
const auto mmod_rects = extract_mmod_rects(truth_image.truth_instances);
cropper(input_image, mmod_rects, temp.input_image, temp.mmod_rects);
disturb_colors(temp.input_image, rnd);
// Push the result to be used by the trainer.
data.enqueue(temp);
}
};
std::thread data_loader1([f]() { f(1); });
std::thread data_loader2([f]() { f(2); });
std::thread data_loader3([f]() { f(3); });
std::thread data_loader4([f]() { f(4); });
const auto stop_data_loaders = [&]()
{
data.disable();
data_loader1.join();
data_loader2.join();
data_loader3.join();
data_loader4.join();
};
dnn_trainer<det_bnet_type> det_trainer(det_net, sgd(weight_decay, momentum));
try
{
det_trainer.be_verbose();
det_trainer.set_learning_rate(initial_learning_rate);
det_trainer.set_synchronization_file("pascal_voc2012_det_trainer_state_file.dat", std::chrono::minutes(10));
det_trainer.set_iterations_without_progress_threshold(5000);
// Output training parameters.
cout << det_trainer << endl;
std::vector<matrix<rgb_pixel>> samples;
std::vector<std::vector<mmod_rect>> labels;
// The main training loop. Keep making mini-batches and giving them to the trainer.
// We will run until the learning rate becomes small enough.
while (det_trainer.get_learning_rate() >= 1e-4)
{
samples.clear();
labels.clear();
// make a mini-batch
det_training_sample temp;
while (samples.size() < det_minibatch_size)
{
data.dequeue(temp);
samples.push_back(std::move(temp.input_image));
labels.push_back(std::move(temp.mmod_rects));
}
det_trainer.train_one_step(samples, labels);
}
}
catch (std::exception&)
{
stop_data_loaders();
throw;
}
// Training done, tell threads to stop and make sure to wait for them to finish before
// moving on.
stop_data_loaders();
// also wait for threaded processing to stop in the trainer.
det_trainer.get_net();
det_net.clean();
return det_net;
}
// ----------------------------------------------------------------------------------------
matrix<float> keep_only_current_instance(const matrix<rgb_pixel>& rgb_label_image, const rgb_pixel rgb_label)
{
const auto nr = rgb_label_image.nr();
const auto nc = rgb_label_image.nc();
matrix<float> result(nr, nc);
for (long r = 0; r < nr; ++r)
{
for (long c = 0; c < nc; ++c)
{
const auto& index = rgb_label_image(r, c);
if (index == rgb_label)
result(r, c) = +1;
else if (index == dlib::rgb_pixel(224, 224, 192))
result(r, c) = 0;
else
result(r, c) = -1;
}
}
return result;
}
seg_bnet_type train_segmentation_network(
const std::vector<truth_image>& truth_images,
unsigned int seg_minibatch_size,
const std::string& classlabel
)
{
seg_bnet_type seg_net;
const double initial_learning_rate = 0.1;
const double weight_decay = 0.0001;
const double momentum = 0.9;
const std::string synchronization_file_name
= "pascal_voc2012_seg_trainer_state_file"
+ (classlabel.empty() ? "" : ("_" + classlabel))
+ ".dat";
dnn_trainer<seg_bnet_type> seg_trainer(seg_net, sgd(weight_decay, momentum));
seg_trainer.be_verbose();
seg_trainer.set_learning_rate(initial_learning_rate);
seg_trainer.set_synchronization_file(synchronization_file_name, std::chrono::minutes(10));
seg_trainer.set_iterations_without_progress_threshold(2000);
set_all_bn_running_stats_window_sizes(seg_net, 1000);
// Output training parameters.
cout << seg_trainer << endl;
std::vector<matrix<rgb_pixel>> samples;
std::vector<matrix<float>> labels;
// Start a bunch of threads that read images from disk and pull out random crops. It's
// important to be sure to feed the GPU fast enough to keep it busy. Using multiple
// thread for this kind of data preparation helps us do that. Each thread puts the
// crops into the data queue.
dlib::pipe<seg_training_sample> data(200);
auto f = [&data, &truth_images](time_t seed)
{
dlib::rand rnd(time(0) + seed);
matrix<rgb_pixel> input_image;
matrix<rgb_pixel> rgb_label_image;
matrix<rgb_pixel> rgb_label_chip;
seg_training_sample temp;
while (data.is_enabled())
{
// Pick a random input image.
const auto random_index = rnd.get_random_32bit_number() % truth_images.size();
const auto& truth_image = truth_images[random_index];
const auto image_truths = truth_image.truth_instances;
if (!image_truths.empty())
{
const image_info& info = truth_image.info;
// Load the input image.
load_image(input_image, info.image_filename);
// Load the ground-truth (RGB) instance labels.
load_image(rgb_label_image, info.instance_label_filename);
// Pick a random training instance.
const auto& truth_instance = image_truths[rnd.get_random_32bit_number() % image_truths.size()];
const auto& truth_rect = truth_instance.mmod_rect.rect;
const auto cropping_rect = get_cropping_rect(truth_rect);
// Pick a random crop around the instance.
const auto max_x_translate_amount = static_cast<long>(truth_rect.width() / 10.0);
const auto max_y_translate_amount = static_cast<long>(truth_rect.height() / 10.0);
const auto random_translate = point(
rnd.get_integer_in_range(-max_x_translate_amount, max_x_translate_amount + 1),
rnd.get_integer_in_range(-max_y_translate_amount, max_y_translate_amount + 1)
);
const rectangle random_rect(
cropping_rect.left() + random_translate.x(),
cropping_rect.top() + random_translate.y(),
cropping_rect.right() + random_translate.x(),
cropping_rect.bottom() + random_translate.y()
);
const chip_details chip_details(random_rect, chip_dims(seg_dim, seg_dim));
// Crop the input image.
extract_image_chip(input_image, chip_details, temp.input_image, interpolate_bilinear());
disturb_colors(temp.input_image, rnd);
// Crop the labels correspondingly. However, note that here bilinear
// interpolation would make absolutely no sense - you wouldn't say that
// a bicycle is half-way between an aeroplane and a bird, would you?
extract_image_chip(rgb_label_image, chip_details, rgb_label_chip, interpolate_nearest_neighbor());
// Clear pixels not related to the current instance.
temp.label_image = keep_only_current_instance(rgb_label_chip, truth_instance.rgb_label);
// Push the result to be used by the trainer.
data.enqueue(temp);
}
else
{
// TODO: use background samples as well
}
}
};
std::thread data_loader1([f]() { f(1); });
std::thread data_loader2([f]() { f(2); });
std::thread data_loader3([f]() { f(3); });
std::thread data_loader4([f]() { f(4); });
const auto stop_data_loaders = [&]()
{
data.disable();
data_loader1.join();
data_loader2.join();
data_loader3.join();
data_loader4.join();
};
try
{
// The main training loop. Keep making mini-batches and giving them to the trainer.
// We will run until the learning rate has dropped by a factor of 1e-4.
while (seg_trainer.get_learning_rate() >= 1e-4)
{
samples.clear();
labels.clear();
// make a mini-batch
seg_training_sample temp;
while (samples.size() < seg_minibatch_size)
{
data.dequeue(temp);
samples.push_back(std::move(temp.input_image));
labels.push_back(std::move(temp.label_image));
}
seg_trainer.train_one_step(samples, labels);
}
}
catch (std::exception&)
{
stop_data_loaders();
throw;
}
// Training done, tell threads to stop and make sure to wait for them to finish before
// moving on.
stop_data_loaders();
// also wait for threaded processing to stop in the trainer.
seg_trainer.get_net();
seg_net.clean();
return seg_net;
}
// ----------------------------------------------------------------------------------------
int ignore_overlapped_boxes(
std::vector<truth_instance>& truth_instances,
const test_box_overlap& overlaps
)
/*!
ensures
- Whenever two rectangles in boxes overlap, according to overlaps(), we set the
smallest box to ignore.
- returns the number of newly ignored boxes.
!*/
{
int num_ignored = 0;
for (size_t i = 0, end = truth_instances.size(); i < end; ++i)
{
auto& box_i = truth_instances[i].mmod_rect;
if (box_i.ignore)
continue;
for (size_t j = i+1; j < end; ++j)
{
auto& box_j = truth_instances[j].mmod_rect;
if (box_j.ignore)
continue;
if (overlaps(box_i, box_j))
{
++num_ignored;
if(box_i.rect.area() < box_j.rect.area())
box_i.ignore = true;
else
box_j.ignore = true;
}
}
}
return num_ignored;
}
std::vector<truth_instance> load_truth_instances(const image_info& info)
{
matrix<rgb_pixel> instance_label_image;
matrix<rgb_pixel> class_label_image;
load_image(instance_label_image, info.instance_label_filename);
load_image(class_label_image, info.class_label_filename);
return rgb_label_images_to_truth_instances(instance_label_image, class_label_image);
}
std::vector<std::vector<truth_instance>> load_all_truth_instances(const std::vector<image_info>& listing)
{
std::vector<std::vector<truth_instance>> truth_instances(listing.size());
parallel_for(
0,
listing.size(),
[&](size_t index)
{
truth_instances[index] = load_truth_instances(listing[index]);
}
);
return truth_instances;
}
// ----------------------------------------------------------------------------------------
std::vector<truth_image> filter_based_on_classlabel(
const std::vector<truth_image>& truth_images,
const std::vector<std::string>& desired_classlabels
)
{
std::vector<truth_image> result;
const auto represents_desired_class = [&desired_classlabels](const truth_instance& truth_instance) {
return std::find(
desired_classlabels.begin(),
desired_classlabels.end(),
truth_instance.mmod_rect.label
) != desired_classlabels.end();
};
for (const auto& input : truth_images)
{
const auto has_desired_class = std::any_of(
input.truth_instances.begin(),
input.truth_instances.end(),
represents_desired_class
);
if (has_desired_class) {
// NB: This keeps only MMOD rects belonging to any of the desired classes.
// A reasonable alternative could be to keep all rects, but mark those
// belonging in other classes to be ignored during training.
std::vector<truth_instance> temp;
std::copy_if(
input.truth_instances.begin(),
input.truth_instances.end(),
std::back_inserter(temp),
represents_desired_class
);
result.push_back(truth_image{ input.info, temp });
}
}
return result;
}
// Ignore truth boxes that overlap too much, are too small, or have a large aspect ratio.
void ignore_some_truth_boxes(std::vector<truth_image>& truth_images)
{
for (auto& i : truth_images)
{
auto& truth_instances = i.truth_instances;
ignore_overlapped_boxes(truth_instances, test_box_overlap(0.90, 0.95));
for (auto& truth : truth_instances)
{
if (truth.mmod_rect.ignore)
continue;
const auto& rect = truth.mmod_rect.rect;
constexpr unsigned long min_width = 35;
constexpr unsigned long min_height = 35;
if (rect.width() < min_width && rect.height() < min_height)
{
truth.mmod_rect.ignore = true;
continue;
}
constexpr double max_aspect_ratio_width_to_height = 3.0;
constexpr double max_aspect_ratio_height_to_width = 1.5;
const double aspect_ratio_width_to_height = rect.width() / static_cast<double>(rect.height());
const double aspect_ratio_height_to_width = 1.0 / aspect_ratio_width_to_height;
const bool is_aspect_ratio_too_large
= aspect_ratio_width_to_height > max_aspect_ratio_width_to_height
|| aspect_ratio_height_to_width > max_aspect_ratio_height_to_width;
if (is_aspect_ratio_too_large)
truth.mmod_rect.ignore = true;
}
}
}
// Filter images that have no (non-ignored) truth
std::vector<truth_image> filter_images_with_no_truth(const std::vector<truth_image>& truth_images)
{
std::vector<truth_image> result;
for (const auto& truth_image : truth_images)
{
const auto ignored = [](const truth_instance& truth) { return truth.mmod_rect.ignore; };
const auto& truth_instances = truth_image.truth_instances;
if (!std::all_of(truth_instances.begin(), truth_instances.end(), ignored))
result.push_back(truth_image);
}
return result;
}
int main(int argc, char** argv) try
{
if (argc < 2)
{
cout << "To run this program you need a copy of the PASCAL VOC2012 dataset." << endl;
cout << endl;
cout << "You call this program like this: " << endl;
cout << "./dnn_instance_segmentation_train_ex /path/to/VOC2012 [det-minibatch-size] [seg-minibatch-size] [class-1] [class-2] [class-3] ..." << endl;
return 1;
}
cout << "\nSCANNING PASCAL VOC2012 DATASET\n" << endl;
const auto listing = get_pascal_voc2012_train_listing(argv[1]);
cout << "images in entire dataset: " << listing.size() << endl;
if (listing.size() == 0)
{
cout << "Didn't find the VOC2012 dataset. " << endl;
return 1;
}
// mini-batches smaller than the default can be used with GPUs having less memory
const unsigned int det_minibatch_size = argc >= 3 ? std::stoi(argv[2]) : 35;
const unsigned int seg_minibatch_size = argc >= 4 ? std::stoi(argv[3]) : 100;
cout << "det mini-batch size: " << det_minibatch_size << endl;
cout << "seg mini-batch size: " << seg_minibatch_size << endl;
std::vector<std::string> desired_classlabels;
for (int arg = 4; arg < argc; ++arg)
desired_classlabels.push_back(argv[arg]);
if (desired_classlabels.empty())
{
desired_classlabels.push_back("bicycle");
desired_classlabels.push_back("car");
desired_classlabels.push_back("cat");
}
cout << "desired classlabels:";
for (const auto& desired_classlabel : desired_classlabels)
cout << " " << desired_classlabel;
cout << endl;
// extract the MMOD rects
cout << endl << "Extracting all truth instances...";
const auto truth_instances = load_all_truth_instances(listing);
cout << " Done!" << endl << endl;
DLIB_CASSERT(listing.size() == truth_instances.size());
std::vector<truth_image> original_truth_images;
for (size_t i = 0, end = listing.size(); i < end; ++i)
{
original_truth_images.push_back(truth_image{
listing[i], truth_instances[i]
});
}
auto truth_images_filtered_by_class = filter_based_on_classlabel(original_truth_images, desired_classlabels);
cout << "images in dataset filtered by class: " << truth_images_filtered_by_class.size() << endl;
ignore_some_truth_boxes(truth_images_filtered_by_class);
const auto truth_images = filter_images_with_no_truth(truth_images_filtered_by_class);
cout << "images in dataset after ignoring some truth boxes: " << truth_images.size() << endl;
// First train an object detector network (loss_mmod).
cout << endl << "Training detector network:" << endl;
const auto det_net = train_detection_network(truth_images, det_minibatch_size);
// Then train mask predictors (segmentation).
std::map<std::string, seg_bnet_type> seg_nets_by_class;
// This flag controls if a separate mask predictor is trained for each class.
// Note that it would also be possible to train a separate mask predictor for
// class groups, each containing somehow similar classes -- for example, one
// mask predictor for cars and buses, another for cats and dogs, and so on.
constexpr bool separate_seg_net_for_each_class = true;
if (separate_seg_net_for_each_class)
{
for (const auto& classlabel : desired_classlabels)
{
// Consider only the truth images belonging to this class.
const auto class_images = filter_based_on_classlabel(truth_images, { classlabel });
cout << endl << "Training segmentation network for class " << classlabel << ":" << endl;
seg_nets_by_class[classlabel] = train_segmentation_network(class_images, seg_minibatch_size, classlabel);
}
}
else
{
cout << "Training a single segmentation network:" << endl;
seg_nets_by_class[""] = train_segmentation_network(truth_images, seg_minibatch_size, "");
}
cout << "Saving networks" << endl;
serialize(instance_segmentation_net_filename) << det_net << seg_nets_by_class;
}
catch(std::exception& e)
{
cout << e.what() << endl;
}