k4abt.hpp

#ifndef K4ABT_HPP
#define K4ABT_HPP

#include <k4a/k4a.hpp>
#include "k4abt.h"

namespace k4abt
{

    class frame
    {
    public:
        frame(k4abt_frame_t handle = nullptr) noexcept : m_handle(handle) {}

        frame(const frame &other) noexcept : m_handle(other.m_handle)
        {
            if (m_handle != nullptr)
            {
                k4abt_frame_reference(m_handle);
            }
        }

        frame(frame &&other) noexcept : m_handle(other.m_handle)
        {
            other.m_handle = nullptr;
        }

        ~frame()
        {
            reset();
        }

        frame &operator=(const frame &other) noexcept
        {
            if (this != &other)
            {
                reset();
                m_handle = other.m_handle;
                if (m_handle != nullptr)
                {
                    k4abt_frame_reference(m_handle);
                }
            }
            return *this;
        }

        frame &operator=(frame &&other) noexcept
        {
            if (this != &other)
            {
                reset();
                m_handle = other.m_handle;
                other.m_handle = nullptr;
            }
            return *this;
        }

        frame &operator=(std::nullptr_t) noexcept
        {
            reset();
            return *this;
        }

        bool operator==(const frame &other) const noexcept
        {
            return m_handle == other.m_handle;
        }

        bool operator==(std::nullptr_t) const noexcept
        {
            return m_handle == nullptr;
        }

        bool operator!=(const frame &other) const noexcept
        {
            return m_handle != other.m_handle;
        }

        bool operator!=(std::nullptr_t) const noexcept
        {
            return m_handle != nullptr;
        }

        operator bool() const noexcept
        {
            return m_handle != nullptr;
        }

        k4abt_frame_t handle() const noexcept
        {
            return m_handle;
        }

        void reset() noexcept
        {
            if (m_handle != nullptr)
            {
                k4abt_frame_release(m_handle);
                m_handle = nullptr;
            }
        }

        uint32_t get_num_bodies() noexcept
        {
            return k4abt_frame_get_num_bodies(m_handle);
        }

        k4abt_skeleton_t get_body_skeleton(uint32_t index)
        {
            k4abt_skeleton_t skeleton;
            get_body_skeleton(index, skeleton);
            return skeleton;
        }

        void get_body_skeleton(uint32_t index, k4abt_skeleton_t& skeleton)
        {
            k4a_result_t result = k4abt_frame_get_body_skeleton(m_handle, index, &skeleton);
            if (K4A_RESULT_SUCCEEDED != result)
            {
                throw k4a::error("Get body skeleton failed!");
            }
        }

        uint32_t get_body_id(uint32_t index) noexcept
        {
            return k4abt_frame_get_body_id(m_handle, index);
        }

        k4abt_body_t get_body(uint32_t index) noexcept
        {
            k4abt_body_t body;
            body.id = k4abt_frame_get_body_id(m_handle, index);
            get_body_skeleton(index, body.skeleton);
            return body;
        }

        std::chrono::microseconds get_device_timestamp() const noexcept
        {
            return std::chrono::microseconds(k4abt_frame_get_device_timestamp_usec(m_handle));
        }

        k4a::image get_body_index_map() const noexcept
        {
            return k4a::image(k4abt_frame_get_body_index_map(m_handle));
        }

        k4a::capture get_capture() const noexcept
        {
            return k4a::capture(k4abt_frame_get_capture(m_handle));
        }

    private:
        k4abt_frame_t m_handle;
    };

    class tracker
    {
    public:
        tracker(k4abt_tracker_t handle = nullptr) noexcept : m_handle(handle) {}

        tracker(tracker &&dev) noexcept : m_handle(dev.m_handle)
        {
            dev.m_handle = nullptr;
        }

        tracker(const tracker &) = delete;

        ~tracker()
        {
            destroy();
        }

        tracker &operator=(const tracker &) = delete;

        tracker &operator=(tracker &&dev) noexcept
        {
            if (this != &dev)
            {
                destroy();
                m_handle = dev.m_handle;
                dev.m_handle = nullptr;
            }
            return *this;
        }

        operator bool() const noexcept
        {
            return m_handle != nullptr;
        }

        void destroy() noexcept
        {
            if (m_handle != nullptr)
            {
                k4abt_tracker_destroy(m_handle);
                m_handle = nullptr;
            }
        }

        bool enqueue_capture(k4a::capture cap, std::chrono::milliseconds timeout = std::chrono::milliseconds(K4A_WAIT_INFINITE))
        {
            int32_t timeout_ms = k4a::internal::clamp_cast<int32_t>(timeout.count());
            k4a_wait_result_t result = k4abt_tracker_enqueue_capture(m_handle, cap.handle(), timeout_ms);
            if (result == K4A_WAIT_RESULT_FAILED)
            {
                throw k4a::error("Failed to enqueue capture to tracker!");
            }
            else if (result == K4A_WAIT_RESULT_TIMEOUT)
            {
                return false;
            }

            return true;
        }

        bool pop_result(k4abt::frame* body_frame, std::chrono::milliseconds timeout = std::chrono::milliseconds(K4A_WAIT_INFINITE))
        {
            k4abt_frame_t frame_handle = nullptr;
            int32_t timeout_ms = k4a::internal::clamp_cast<int32_t>(timeout.count());
            k4a_wait_result_t result = k4abt_tracker_pop_result(m_handle, &frame_handle, timeout_ms);
            if (result == K4A_WAIT_RESULT_FAILED)
            {
                throw k4a::error("Failed to pop result from tracker!");
            }
            else if (result == K4A_WAIT_RESULT_TIMEOUT)
            {
                return false;
            }

            *body_frame = frame(frame_handle);
            return true;
        }

        k4abt::frame pop_result(std::chrono::milliseconds timeout = std::chrono::milliseconds(K4A_WAIT_INFINITE))
        {
            k4abt::frame body_frame = nullptr;
            if (pop_result(&body_frame, timeout))
            {
                return body_frame;
            }
            else
            {
                return nullptr;
            }
        }

        void set_temporal_smoothing(float smoothing_factor) noexcept
        {
            k4abt_tracker_set_temporal_smoothing(m_handle, smoothing_factor);
        }

        void shutdown() noexcept
        {
            k4abt_tracker_shutdown(m_handle);
        }

        static tracker create(const k4a_calibration_t& sensor_calibration, const k4abt_tracker_configuration_t& config = K4ABT_TRACKER_CONFIG_DEFAULT)
        {
            k4abt_tracker_t handle = nullptr;
            k4a_result_t result = k4abt_tracker_create(&sensor_calibration, config, &handle);

            if (K4A_RESULT_SUCCEEDED != result)
            {
                throw k4a::error("Failed to create k4abt tracker!");
            }
            return tracker(handle);
        }

    private:
        k4abt_tracker_t m_handle;
    };

} // namespace k4abt

#endif