image_meta_consumer.h

Go to the documentation of this file.

/*
 * Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */
 
#pragma once
 
#include <condition_variable>
#include <mutex>
#include <string>
#include <atomic>
#include <thread>
#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <ios>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <deepstream_config.h>
#include "gstnvdsmeta.h"
#include "nvbufsurface.h"
#include "gst-nvmessage.h"
#include "nvds_obj_encode.h"
#include "concurrent_queue.h"
#include "capture_time_rules.h"
 
class ImageMetaConsumer {
public:
    enum OutputType {
        KITTI = 0,
        JSON = 1,
        CSV = 2
    };
 
    enum ImageSizeType {
        FULL_FRAME,
        CROPPED_TO_OBJECT
    };
 
    ImageMetaConsumer();
 
    ~ImageMetaConsumer();
 
    void init(unsigned gpu_id, const std::string &output_folder_path, const std::string &frame_to_skip_rules_path,
              float min_box_confidence, float max_box_confidence,
              unsigned min_box_width, unsigned min_box_height,
              bool save_full_frame_enabled, bool save_cropped_obj_enabled,
              unsigned seconds_to_skip_interval, unsigned source_nb);
 
    void add_meta_csv(const std::string &meta);
 
    void add_meta_json(const std::string &meta);
 
    void add_meta_kitti(const std::pair<std::string, std::string> &meta);
 
    void stop();
 
 
    float get_min_confidence() const;
 
    float get_max_confidence() const;
 
    unsigned get_gpu_id() const;
 
    unsigned get_min_box_width() const;
 
    unsigned get_min_box_height() const;
 
    bool get_is_stopped() const;
 
    bool get_save_full_frame_enabled() const;
 
    bool get_save_cropped_images_enabled() const;
 
    NvDsObjEncCtxHandle get_obj_ctx_handle();
 
    std::string make_img_path(ImageMetaConsumer::ImageSizeType ist,
                              unsigned stream_source_id,
                              const std::string &datetime_iso8601);
 
    void init_image_save_library_on_first_time();
 
    bool should_save_data(unsigned source_id);
 
    void lock_source_nb(unsigned source_id);
 
    void data_was_saved_for_source(unsigned source_id);
 
    void unlock_source_nb(unsigned source_id);
 
private:
 
    void run();
 
    void multi_metadata_maker(ConcurrentQueue<std::pair<std::string, std::string>> &queue,
                              std::mutex &mutex, std::condition_variable &cv);
 
    bool setup_files();
 
    bool setup_folders();
 
    void write_intro(std::ofstream &os, OutputType &ot);
 
    void write_mid_separator(std::ofstream &os, OutputType &ot);
 
    void write_end(std::ofstream &os, OutputType &ot, unsigned total_nb);
 
    unsigned int get_unique_id();
 
    void single_metadata_maker(const std::string &extension,
                               ConcurrentQueue<std::string> &queue,
                               std::mutex &mutex, std::condition_variable &cv,
                               OutputType ot);
 
    ConcurrentQueue<std::pair<std::string, std::string>> queue_kitti_;
    ConcurrentQueue<std::string> queue_csv_;
    ConcurrentQueue<std::string> queue_json_;
    std::mutex m_kitti_;
    std::mutex m_json_;
    std::mutex m_csv_;
    std::condition_variable cv_kitti_;
    std::condition_variable cv_json_;
    std::condition_variable cv_csv_;
    std::atomic<bool> is_stopped_;
    std::string output_folder_path_;
    std::string images_cropped_obj_output_folder_;
    std::string images_full_frame_output_folder_;
    std::string labels_output_folder_;
    std::thread th_kitti_;
    std::thread th_json_;
    std::thread th_csv_;
    std::mutex mutex_unique_index_;
    unsigned int unique_index_ = 0;
    float min_confidence_;
    float max_confidence_;
    unsigned gpu_id_;
    unsigned min_box_width_;
    unsigned min_box_height_;
    bool save_full_frame_enabled_;
    bool save_cropped_obj_enabled_;
    std::vector<std::chrono::_V2::system_clock::time_point> time_last_frame_saved_list_;
    std::array<std::mutex, MAX_SOURCE_BINS> mutex_frame_saved_list_;
    CaptureTimeRules ctr_;
    NvDsObjEncCtxHandle obj_ctx_handle_;
    bool image_saving_library_is_init_;
};