    // 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;
    }