Vzense-Tof-3D/vzense/device_tof3d.h

#pragma once

#include <thread>
#include <boost/chrono/duration.hpp>
#include <opencv2/opencv.hpp>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>

#include "error_code/error_code.hpp"
#include "VzenseNebula_api.h"
#include "log/log.h"
#include "thread/thread_condition.h"
#include "stream/streamServer.h"
#include "tool/timedlockdata.hpp"
#include "tool/ip.hpp"
#include "pcl/trans_coor.hpp"
#include "pcl/point3D_tool.h"
#include "tof3d_config.pb.h"

#ifdef ENABLE_TENSORRT_DETECT
#include "detect/tensorrt/wheel-detector.h"
#else

#include "detect/onnx/wheel-detector.h"

#endif


class DeviceTof3D {
public:
    struct WheelDetectResult {
        bool haveCloud = false;
        bool segmentOk = false;
        bool detectOk = false;
        pcl::PointXYZ center;
        float confidence = 0;

        void reset() {
            haveCloud = false;
            segmentOk = false;
            detectOk = false;
            center.x = 0;
            center.y = 0;
            center.z = 0;
            confidence = 0;
        }
    };
    struct CarDetectResult {
        float wheels_center_x;
        float wheels_center_y;
        float width;
        float wheel_width;
        float length;
        float wheel_length;
        float front_wheels_theta;
        float theta;
        float wheel_base;

        void reset() {
            wheels_center_x = 0;
            wheels_center_y = 0;
            width = 0;
            wheel_width = 0;
            length = 0;
            wheel_length = 0;
            front_wheels_theta = 0;
            theta = 0;
            wheel_base = 0;
        }
    };

    struct DetectResult {
        WheelDetectResult wheel[DeviceAzimuth::MAX];
        CarDetectResult car;

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

    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;

    // TODO
    void run();

    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(pcl::PointCloud<pcl::PointXYZ>::Ptr &wheel_cloud, DetectResult &result);

    // 车身点云优化
    Error_manager CarPoseOptimize(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 pointMat;

        DeviceMat() {
            depthMat.create(480, 640, CV_16UC1);
            pointMat.create(480, 640, CV_32FC4);
        }

        DeviceMat& operator=(const DeviceMat &p) {
            this->depthMat = p.depthMat.clone();
            this->pointMat = p.pointMat.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;
};