gsd/framework/request/include/CNStreamInferData.h

#pragma once

#include <iostream>
#include <vector>
#include <rapidjson/document.h>
#include <rapidjson/rapidjson.h>
#include <rapidjson/stringbuffer.h>
#include <rapidjson/writer.h>

using namespace std;

typedef struct {
  double x, y, w, h;
} InferBoundingBox;

class InferInfo
{
public:
    std::string Label;
    double Score;
    InferBoundingBox BBox;

public:
    InferInfo() {}
    ~InferInfo() {}

    bool jsonToObject(const rapidjson::Value& object){
        const auto end = object.MemberEnd();

        if(end == object.FindMember("Label") || !object["Label"].IsString()){
            return false;
        }
        else{
            Label = object["Label"].GetString();
        }

        if(end == object.FindMember("Score") || !object["Score"].IsDouble()){
            return false;
        }
        else{
            Score = object["Score"].GetDouble();
        }

        if(end == object.FindMember("BBox") || !object["BBox"].IsArray()){
            return false;
        }
        else{
            const rapidjson::Value& objs = object["BBox"];
            BBox.x = objs[0].GetDouble();
            BBox.y = objs[1].GetDouble();
            BBox.w = objs[2].GetDouble();
            BBox.h = objs[3].GetDouble();
        }
        return true;
    }
};


class CNStreamInferData
{
public:
    using Ptr = std::shared_ptr<CNStreamInferData>;
    
    std::string StreamName;
    int FrameCount;
    int width;
    int height;

    vector<InferInfo> Objects;
    
    std::string ImageBase64;
    std::string videoPath;

public:
    CNStreamInferData() {}
    ~CNStreamInferData() {}

    bool jsonToObject(const std::string& json){
        rapidjson::Document doc;
        if (doc.Parse<rapidjson::kParseCommentsFlag>(json.c_str()).HasParseError()) {
            return false;
        }
        // get members
        const auto end = doc.MemberEnd();

        // json_type
        if (end == doc.FindMember("StreamName") || !doc["StreamName"].IsString()) {
            return false;
        } else {
            StreamName = doc["StreamName"].GetString();
        }

        if (end == doc.FindMember("width") || !doc["width"].IsInt()) {
            return false;
        } else {
            width = doc["width"].GetInt();
        }

        if (end == doc.FindMember("height") || !doc["height"].IsInt()) {
            return false;
        } else {
            height = doc["height"].GetInt();
        }

        if (end == doc.FindMember("FrameCount") || !doc["FrameCount"].IsInt()) {
            return false;
        } else {
            FrameCount = doc["FrameCount"].GetInt();
        }

        if (end == doc.FindMember("Objects") || !doc["Objects"].IsArray()) {
            return false;
        } else {
            const rapidjson::Value& objects = doc["Objects"];
            for (size_t i = 0; i < objects.Size(); ++i) {
                const rapidjson::Value& obj =  objects[i];
                InferInfo data;
                if(!data.jsonToObject(obj)) return false;
                this->Objects.push_back(data);
            }
        }

        if(end != doc.FindMember("ImageBase64") && doc["ImageBase64"].IsString())
        {
            ImageBase64 = doc["ImageBase64"].GetString();
        }

        if (end != doc.FindMember("videoPath") && doc["videoPath"].IsString()) {
            videoPath = doc["videoPath"].GetString();
        } 
        
        return true;
    }
};