image_meta_consumer.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */
 
#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_;
    std::mutex mutex_image_save_init_;
};