Vzense-Tof-3D/vzense/back/device_tof3d_b.h

#pragma once

#include <thread>
#include <boost/chrono/duration.hpp>
#include <opencv2/opencv.hpp>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/approximate_voxel_grid.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/common/transforms.h>
#include <pcl/visualization/pcl_visualizer.h>

#include "error_code/error_code.hpp"
#include "VzenseNebula_api.h"
#include "log/log.h"
#include "thread/thread_condition.h"
#include "communication/grpc/streamServer.h"
#include "tool/timedlockdata.hpp"
#include "tool/ip.hpp"
#include "pcl/trans_coor.hpp"
#include "pcl/point3D_tool.h"
#include "proto/vzense.pb.h"

#ifdef ENABLE_TENSORRT_DETECT
#include "detect/tensorrt_detect/wheel-detector.h"
#else
#include "detect/onnx_detect/wheel-detector.h"
#endif
#include "detect/ceres_detect/car_pose_detector.h"
#include "detect/ceres_detect/wheel_detector.h"
#include "detect/ceres_detect/detect_wheel_ceres.h"

class DeviceTof3D {
public:
    struct DetectResult {
        WheelCeresDetector::WheelDetectResult wheel[DeviceAzimuth::MAX];
        Car_pose_detector::CarDetectResult car;

        void reset() {
            for (int i = 0; i < DeviceAzimuth::MAX; i++) {
                wheel[i].reset();
            }
            car.reset();
        }

        std::string info() {
            std::string str = car.info();
            return str;
        }
    };

    struct tof3dVzenseInfo {
        VzDeviceHandle handle = nullptr;
        bool is_connect = false;
        bool is_start_stream = false;
        VzDeviceInfo info;
        tof3dVzenseEtc etc;
    };

    using VzEtcMap = std::map<std::string, tof3dVzenseEtc>;
    using tof3dVzenseInfoMap = std::map<std::string, tof3dVzenseInfo *>;
public:
    static DeviceTof3D *iter() {
        static DeviceTof3D *instance = nullptr;
        if (instance == nullptr) {
            instance = new DeviceTof3D();
        }
        return instance;
    }

    ~DeviceTof3D() = default;

    Error_manager Init(const VzEtcMap &tp_tof3d, bool search = true);

    Error_manager reInit(const VzEtcMap &etc);

    Error_manager SearchDevice(const double &time);

    Error_manager ConnectDevice(const std::string &ip, bool open_stream = true);

    Error_manager ConnectAllDevices(bool open_stream = true);

    Error_manager ConnectAllEtcDevices(bool open_stream = true);

    Error_manager DisConnectDevice(const std::string &ip);

    Error_manager DisConnectAllEtcDevices();

    Error_manager DisConnectAllDevices();

    Error_manager getDepthFrame(const std::string &ip, VzFrame &depthFrame);

    Error_manager getIrFrame(const std::string &ip, VzFrame &irFrame);

    Error_manager getDepthAndIrPicture(const std::string &ip, VzFrame &depthFrame, VzFrame &irFrame);

    Error_manager getDepthPointCloud(const std::string &ip, pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud);

    Error_manager
    DepthFrame2PointCloud(const std::string &ip, VzFrame &depthFrame, pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud);

    Error_manager updateTof3dEtc();

    Error_manager setTof3dParams(const std::string &ip);

    Error_manager getTof3dParams(const std::string &ip);

    static Error_manager Frame2Mat(VzFrame &frame, cv::Mat &mat);

protected:
    DeviceTof3D() = default;

    static Error_manager DepthFrame2Mat(VzFrame &frame, cv::Mat &mat);

    static Error_manager IrFrame2Mat(VzFrame &frame, cv::Mat &mat);

    void receiveFrameThread(const std::string &ip);

    void detectThread();

    static void HotPlugStateCallback(const VzDeviceInfo *pInfo, int status, void *contex);

    bool drawBox(cv::Mat &mat, cv::Rect &box, cv::Scalar &color, std::string className, float confidence);

    Error_manager loadEtc(const VzEtcMap &etc);

    void stopWorking();

    // // 传入深度图，得到车轮在mat中的坐标
    // Error_manager segmentTensorRT(cv::Mat &depth_mat, std::vector<cv::Point> &wheel_cv_cloud);

    // 根据传入坐标，分离depth frame的数据点
    Error_manager segFrame2CarAndWheel(VzFrame &depth_frame, std::vector<cv::Point> &wheel_cv_cloud,
                                       pcl::PointCloud<pcl::PointXYZ>::Ptr &car_cloud,
                                       pcl::PointCloud<pcl::PointXYZ>::Ptr &wheel_cloud);

    // 车轮点云优化
    Error_manager WheelCloudOptimize(std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> &vct_wheel_cloud, DetectResult &result);

    // 车身点云优化
    Error_manager CarPoseOptimize(pcl::visualization::PCLVisualizer &viewer, int *view_port, pcl::PointCloud<pcl::PointXYZ>::Ptr &car_cloud, DetectResult &result);

    Error_manager grpcVzenseMat(JetStream::ResFrame &frame, const DeviceAzimuth &azimuth, cv::Mat &mat);

    Error_manager grpcVzensePointMat(JetStream::ResFrame &frame, const DeviceAzimuth &azimuth, cv::Mat &mat);

    Error_manager grpcCloud(JetStream::ResFrame &frame, const DeviceAzimuth &azimuth, pcl::PointCloud<pcl::PointXYZI>::Ptr &cloud);

private:
    struct ThreadInfo {
        std::thread *t;
        Thread_condition *condit;

        ThreadInfo() {
            t = nullptr;
            condit = nullptr;
        }

        ThreadInfo(std::thread *m_t, Thread_condition *m_condit) : t(m_t), condit(m_condit) {}
    };

    struct DeviceMat {
        cv::Mat depthMat;
        cv::Mat irMat;
        cv::Mat pointMat;

        DeviceMat() {
            irMat = cv::Mat::zeros(480, 640, CV_8UC1);
            depthMat = cv::Mat::zeros(480, 640, CV_16UC1);
            pointMat = cv::Mat::zeros(480, 640, CV_32FC4);
        }

        DeviceMat& operator=(const DeviceMat &p) {
            this->depthMat = p.depthMat.clone();
            this->pointMat = p.pointMat.clone();
            this->irMat = p.irMat.clone();
            return *this;
        }
    };
    TimedLockData<DeviceMat> m_device_mat[DeviceAzimuth::MAX];

    VzReturnStatus m_vz_status = VzRetOthers;

    tof3dVzenseInfoMap mp_device_info;
    std::map<std::string, ThreadInfo> mp_thread_info;

    ThreadInfo detect_thread;
    TensorrtWheelDetector *detector;
    StreamRpcServer m_grpc_server;

    DetectResult m_detect_result_record;
};