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;
    }


    // std::string Label;
    // double Score;
    // InferBoundingBox BBox;

    /**
     * @description: objectToJson
     * @return {*}
     */    
    void objectToJson(rapidjson::Writer<rapidjson::StringBuffer>& writer){
        writer.StartObject();
        writer.Key("Label");
        writer.String(Label.c_str());

        writer.Key("Score");
        writer.Double(Score);

        writer.Key("BBox");
        writer.StartArray();

        writer.Double(BBox.x);
        writer.Double(BBox.y);
        writer.Double(BBox.w);
        writer.Double(BBox.h);

        writer.EndArray();

        writer.EndObject();
    }
};


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;
        doc.Parse(json.c_str());
        if(!doc.IsObject()){
            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;
    }


    /**
     * @description: objectToJson
     * @param {string&} str
     * @return {*}
     */    
    void objectToJson(std::string& str){
        rapidjson::StringBuffer strBuf;
        rapidjson::Writer<rapidjson::StringBuffer> writer(strBuf);
        this->objectToJson(writer);
        str = strBuf.GetString();
    }


    
    // std::string StreamName;
    // int FrameCount;
    // int width;
    // int height;

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

    /**
     * @description: objectToJson
     * @return {*}
     */    
    void objectToJson(rapidjson::Writer<rapidjson::StringBuffer>& writer){
        writer.StartObject();
        writer.Key("StreamName");
        writer.String(StreamName.c_str());

        writer.Key("FrameCount");
        writer.Int(FrameCount);

        writer.Key("width");
        writer.Int(width);

        writer.Key("height");
        writer.Int(height);

        writer.Key("Objects");
        writer.StartArray();
        for(auto iter = Objects.begin(); iter != Objects.end(); iter++){
            iter->objectToJson(writer);
        }
        writer.EndArray();

        writer.Key("videoPath");
        writer.String(videoPath.c_str());

        writer.EndObject();
    }
};