GBMOT/util/Visualizer.cpp

//
// Created by wrede on 04.05.16.
//

#include "Visualizer.h"
#include "Logger.h"
#include "../core/ObjectDataAngular.h"
#include "FileIO.h"

namespace util
{
    int Visualizer::GetTime()
    {
        return std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::system_clock::now().time_since_epoch()).count();
    }

    void Visualizer::Display(std::vector<core::TrackletPtr> & tracks, size_t frame_offset,
                             std::string image_folder, bool output, bool display,
                             std::string output_path, std::string title, size_t first_frame,
                             int play_fps, bool show_grid, int grid_width, int grid_height)
    {
        util::Logger::LogInfo("Displaying data");

        size_t current_frame = first_frame;

        // Load images
        std::vector<std::string> image_files;
        util::FileIO::ListFiles(image_folder, image_files);

        // Generate a random color for each individual track
        std::vector<cv::Scalar> colors;
        std::random_device rd;
        std::mt19937 gen(rd());
        //TODO find better colors (more colors)
        colors.push_back(cv::Scalar(0, 255, 0));
        colors.push_back(cv::Scalar(0, 0, 255));
        colors.push_back(cv::Scalar(0, 255, 255));
        colors.push_back(cv::Scalar(255, 0, 0));
        for (size_t i = 4; i < tracks.size(); ++i)
        {
            // BGR
            cv::Scalar color(std::generate_canonical<double, 10>(gen) * 255,
                             std::generate_canonical<double, 10>(gen) * 255,
                             std::generate_canonical<double, 10>(gen) * 255);
            colors.push_back(color);
        }

        //TODO move to extra class
        //TODO create necessary directories
        if (output)
        {
            util::Logger::LogInfo("Start writing output images");
            for (size_t i = 0; i < image_files.size(); ++i)
            {
                cv::Mat image = cv::imread(image_folder + "/" + image_files[i], 1);
                for (size_t j = 0; j < tracks.size(); ++j)
                {
                    tracks[j]->Visualize(image, colors[j], frame_offset + i, 0, 0);
                }
                cv::imwrite(output_path + "/images/" + image_files[i], image);
            }
            util::Logger::LogInfo("Finished writing output images");
        }

        if (display)
        {
            // Create window
            cv::namedWindow(title, CV_WINDOW_AUTOSIZE);

            // Display frames and data
            int target_delay = 1000 / play_fps;
            int last_frame_time = GetTime();
            int current_delay, current_frame_time;
            bool play = false;
            while (true)
            {
                // Display image
                cv::Mat image = cv::imread(image_folder + "/" + image_files[current_frame], 1);

                // Draw grid if resolution is specified
                if (show_grid && grid_width > 0 && grid_height > 0)
                {
                    cv::Scalar gridColor(255, 255, 255);
                    double cellWidth = image.cols / grid_width;
                    double cellHeight = image.rows / grid_height;
                    for (int x = 0; x < grid_width; ++x)
                    {
                        for (int y = 0; y < grid_height; ++y)
                        {
                            cv::rectangle(image,
                                          cv::Point2d(x * cellWidth, y * cellHeight),
                                          cv::Point2d((x + 1) * cellWidth, (y + 1) * cellHeight),
                                          gridColor);
                        }
                    }
                }

                // Visualize the tracklets
                for (size_t i = 0; i < tracks.size(); ++i)
                {
                    tracks[i]->Visualize(image, colors[i], current_frame + frame_offset, 0, 0);
                }

                cv::imshow(title, image);

                // Get key input
                int key;
                if (play)
                {
                    current_frame_time = GetTime();
                    current_delay = last_frame_time - current_frame_time;
                    if (current_delay < target_delay)
                    {
                        key = cv::waitKey(target_delay - current_delay);
                    }
                    else
                    {
                        key = 0;
                    }
                    last_frame_time = GetTime();
                }
                else
                {
                    key = cv::waitKey(0);
                }

                // Process key input
                if (key == 1048678) // F
                {
                    play = !play;
                }
                else if (key == 1048603) // ESC
                {
                    break;
                }

                if (play || key == 1048676) // D
                {
                    if (current_frame < image_files.size() - 1)
                    {
                        current_frame++;
                    }
                    else
                    {
                        current_frame = image_files.size() - 1;
                        play = false;
                    }
                }
                else if (key == 1048673) // A
                {
                    if (current_frame > 0)
                    {
                        current_frame--;
                    }
                }
            }
        }
    }
}