Embedded_AI/npu/eeasy/include/Detection.hpp

#pragma once

#include <cstdint>
#include <vector>
#include <algorithm>
#include <cmath>
#include <string>
#include <float.h>

#include <iostream>
#include "AIStatus.h"

#ifndef CLIPRETINA
#define CLIPRETINA(v, mn, mx) \
    {                         \
        if ((v) < (mn))       \
        {                     \
            (v) = (mn);       \
        }                     \
        else if ((v) > (mx))  \
        {                     \
            (v) = (mx);       \
        }                     \
    }
#endif

namespace detection
{
    // sigmoid
    static float sigmoid(float x)
    {

        return 1.0 / (1.0 + exp(-x));
    }

    struct Box
    {
        float xyxy[4] = {0, 0, 0, 0};
        float object_score = 0;
        size_t index = 0;
        float score = 0;
        float area = 0;
    };

    static bool yolov_box_cmp(const Box &a, const Box &b)
    {
        return a.score > b.score;
    }

    static void yolo_nms(std::vector<Box> &boxes, const float &nms_threshold)
    {
        std::sort(boxes.begin(), boxes.end(), yolov_box_cmp);

        size_t current_index = 0;
        while (current_index < boxes.size())
        {
            Box current_box = boxes[current_index];
            size_t running_index = current_index + 1;
            while (running_index < boxes.size())
            {
                Box running_box = boxes[running_index];

                float xx1 = std::max(current_box.xyxy[0], running_box.xyxy[0]);
                float yy1 = std::max(current_box.xyxy[1], running_box.xyxy[1]);
                float xx2 = std::min(current_box.xyxy[2], running_box.xyxy[2]);
                float yy2 = std::min(current_box.xyxy[3], running_box.xyxy[3]);
                float w = std::max(0.0f, xx2 - xx1 + 1.0f);
                float h = std::max(0.0f, yy2 - yy1 + 1.0f);

                float inter_area = w * h;
                float union_area = current_box.area + running_box.area - inter_area;
                float overlap = inter_area / union_area;
                if (overlap > nms_threshold)
                {
                    boxes.erase(boxes.begin() + running_index);
                }
                else
                {
                    ++running_index;
                }
            }
            ++current_index;
        }
    }

    static void generate_proposals_yolov5(std::vector<Box> &proposals,
                                          const signed char *output_ptr,
                                          std::array<vx_size, 4U> output_size,
                                          const unsigned int &stride,
                                          const int &fl,
                                          const int &data_type,
                                          const unsigned int *anchor,
                                          const std::vector<std::string> class_names,
                                          std::unordered_map<std::string, float> class_attributes,
                                          float lookup_table[][6])
    {
        // printf("[%s Line:%d] dim:[%d %d %d %d] anchor:[%d,%d %d,%d %d,%d] stride:%d fl:%d data_type:%d class_num:%d\n", __FUNCTION__, __LINE__,

        //    output_size[0], output_size[1], output_size[2], output_size[3],
        //    anchor[0], anchor[1], anchor[2], anchor[3], anchor[4], anchor[5],
        //    stride, fl, data_type, class_names.size());
        int H_algin = 0;
        if (data_type == 2)
            H_algin = (output_size[1] + 3) / 4 * 4;
        else
            H_algin = (output_size[1] + 1) / 2 * 2;

        int class_num = class_names.size();
        // #pragma unroll
        for (size_t a = 0; a < 3; ++a)
        { // anchor groups = 3
            for (size_t w = 0; w < output_size[0]; ++w)
            {
                for (size_t h = 0; h < output_size[1]; ++h)
                {
                    Box box;
                    size_t max_index = 0;
                    float max_score = -1;
                    // #pragma unroll
                    // for (size_t c = 0; c < 4 + 1 + class_num; ++c)
                    for (size_t c = 0; c < static_cast<size_t>(4 + 1 + class_num); ++c)
                    {
                        size_t ci = a * (4 + 1 + class_num) + c;
                        size_t index = ci / 16 * output_size[0] * H_algin * 16 + w * H_algin * 16 + h * 16 + (ci % 16);

                        // scale and sigmoid
                        // float data = sigmoid(output_ptr[index] * 1.0 / pow(2, fl));
                        float data = lookup_table[output_ptr[index] - (-128)][fl];
                        if (c == 0)
                        {
                            data = (data * 2 - 0.5f + w) * static_cast<float>(stride);
                            box.xyxy[c] = data;
                        }
                        else if (c == 1)
                        {
                            data = (data * 2 - 0.5f + h) * static_cast<float>(stride);
                            box.xyxy[c] = data;
                        }
                        else if (c == 2 || c == 3)
                        {
                            data = powf((data * 2), 2) * anchor[a * 2 + c - 2];
                            box.xyxy[c] = data;
                        }
                        else if (c == 4)
                        {
                            box.object_score = data;
                        }
                        else
                        {
                            if (data > max_score)
                            {
                                max_index = c - 5;
                                max_score = data;
                            }
                        }
                    }
                
                    box.score = max_score * box.object_score;
                    box.index = max_index;

                    bool is_push = false;

                    if (class_attributes.find(class_names[box.index]) != class_attributes.end())
                    {

                        if (box.object_score > class_attributes.find(class_names[box.index])->second &&
                            box.score > class_attributes.find(class_names[box.index])->second)
                        {
                            is_push = true;
                        }
                    }
                    else
                    {
                        if (box.object_score > class_attributes.find("all")->second &&
                            box.score > class_attributes.find("all")->second)
                        {
                            is_push = true;
                        }
                    }

                    // if (box.object_score > 0.5 && box.score > 0.5)
                    // {
                    //     is_push = true;
                    // }

                    if (is_push)
                    {
                        // xywh -> xyxy
                        float x = box.xyxy[0], y = box.xyxy[1], w = box.xyxy[2], h = box.xyxy[3];
                        box.xyxy[0] = x - w / 2;
                        box.xyxy[1] = y - h / 2;
                        box.xyxy[2] = x + w / 2;
                        box.xyxy[3] = y + h / 2;
                        box.area = (box.xyxy[2] - box.xyxy[0] + 1) * (box.xyxy[3] - box.xyxy[1] + 1);
                        proposals.push_back(box);
                        // printf("[%s Line:%d] ------------box %f %f %f %f %f %d------------\n", __FUNCTION__, __LINE__, box.xyxy[0], box.xyxy[1], box.xyxy[2], box.xyxy[3], box.score, box.index);
                    }
                }
            }
        }
    }

}