CNStreamInferData.h 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. #pragma once
  2. #include <iostream>
  3. #include <vector>
  4. #include <rapidjson/document.h>
  5. #include <rapidjson/rapidjson.h>
  6. #include <rapidjson/stringbuffer.h>
  7. #include <rapidjson/writer.h>
  8. using namespace std;
  9. typedef struct {
  10. double x, y, w, h;
  11. } InferBoundingBox;
  12. class InferInfo
  13. {
  14. public:
  15. std::string Label;
  16. double Score;
  17. InferBoundingBox BBox;
  18. public:
  19. InferInfo() {}
  20. ~InferInfo() {}
  21. bool jsonToObject(const rapidjson::Value& object){
  22. const auto end = object.MemberEnd();
  23. if(end == object.FindMember("Label") || !object["Label"].IsString()){
  24. return false;
  25. }
  26. else{
  27. Label = object["Label"].GetString();
  28. }
  29. if(end == object.FindMember("Score") || !object["Score"].IsDouble()){
  30. return false;
  31. }
  32. else{
  33. Score = object["Score"].GetDouble();
  34. }
  35. if(end == object.FindMember("BBox") || !object["BBox"].IsArray()){
  36. return false;
  37. }
  38. else{
  39. const rapidjson::Value& objs = object["BBox"];
  40. BBox.x = objs[0].GetDouble();
  41. BBox.y = objs[1].GetDouble();
  42. BBox.w = objs[2].GetDouble();
  43. BBox.h = objs[3].GetDouble();
  44. }
  45. return true;
  46. }
  47. // std::string Label;
  48. // double Score;
  49. // InferBoundingBox BBox;
  50. /**
  51. * @description: objectToJson
  52. * @return {*}
  53. */
  54. void objectToJson(rapidjson::Writer<rapidjson::StringBuffer>& writer){
  55. writer.StartObject();
  56. writer.Key("Label");
  57. writer.String(Label.c_str());
  58. writer.Key("Score");
  59. writer.Double(Score);
  60. writer.Key("BBox");
  61. writer.StartArray();
  62. writer.Double(BBox.x);
  63. writer.Double(BBox.y);
  64. writer.Double(BBox.w);
  65. writer.Double(BBox.h);
  66. writer.EndArray();
  67. writer.EndObject();
  68. }
  69. };
  70. class CNStreamInferData
  71. {
  72. public:
  73. using Ptr = std::shared_ptr<CNStreamInferData>;
  74. std::string StreamName;
  75. int FrameCount = 0;
  76. int width;
  77. int height;
  78. vector<InferInfo> Objects;
  79. std::string ImageBase64;
  80. std::string videoPath;
  81. public:
  82. CNStreamInferData() {}
  83. ~CNStreamInferData() {}
  84. bool jsonToObject(const std::string& json){
  85. rapidjson::Document doc;
  86. doc.Parse(json.c_str());
  87. if(!doc.IsObject()){
  88. return false;
  89. }
  90. // get members
  91. const auto end = doc.MemberEnd();
  92. // json_type
  93. if (end == doc.FindMember("StreamName") || !doc["StreamName"].IsString()) {
  94. return false;
  95. } else {
  96. StreamName = doc["StreamName"].GetString();
  97. }
  98. if (end == doc.FindMember("width") || !doc["width"].IsInt()) {
  99. return false;
  100. } else {
  101. width = doc["width"].GetInt();
  102. }
  103. if (end == doc.FindMember("height") || !doc["height"].IsInt()) {
  104. return false;
  105. } else {
  106. height = doc["height"].GetInt();
  107. }
  108. if (end == doc.FindMember("FrameCount") || !doc["FrameCount"].IsInt()) {
  109. return false;
  110. } else {
  111. FrameCount = doc["FrameCount"].GetInt();
  112. }
  113. if (end == doc.FindMember("Objects") || !doc["Objects"].IsArray()) {
  114. return false;
  115. } else {
  116. const rapidjson::Value& objects = doc["Objects"];
  117. for (size_t i = 0; i < objects.Size(); ++i) {
  118. const rapidjson::Value& obj = objects[i];
  119. InferInfo data;
  120. if(!data.jsonToObject(obj)) return false;
  121. this->Objects.push_back(data);
  122. }
  123. }
  124. if(end != doc.FindMember("ImageBase64") && doc["ImageBase64"].IsString())
  125. {
  126. ImageBase64 = doc["ImageBase64"].GetString();
  127. }
  128. if (end != doc.FindMember("videoPath") && doc["videoPath"].IsString()) {
  129. videoPath = doc["videoPath"].GetString();
  130. }
  131. return true;
  132. }
  133. /**
  134. * @description: objectToJson
  135. * @param {string&} str
  136. * @return {*}
  137. */
  138. void objectToJson(std::string& str){
  139. rapidjson::StringBuffer strBuf;
  140. rapidjson::Writer<rapidjson::StringBuffer> writer(strBuf);
  141. this->objectToJson(writer);
  142. str = strBuf.GetString();
  143. }
  144. // std::string StreamName;
  145. // int FrameCount;
  146. // int width;
  147. // int height;
  148. // vector<InferInfo> Objects;
  149. // std::string ImageBase64;
  150. // std::string videoPath;
  151. /**
  152. * @description: objectToJson
  153. * @return {*}
  154. */
  155. void objectToJson(rapidjson::Writer<rapidjson::StringBuffer>& writer){
  156. writer.StartObject();
  157. writer.Key("StreamName");
  158. writer.String(StreamName.c_str());
  159. writer.Key("FrameCount");
  160. writer.Int(FrameCount);
  161. writer.Key("width");
  162. writer.Int(width);
  163. writer.Key("height");
  164. writer.Int(height);
  165. writer.Key("Objects");
  166. writer.StartArray();
  167. for(auto iter = Objects.begin(); iter != Objects.end(); iter++){
  168. iter->objectToJson(writer);
  169. }
  170. writer.EndArray();
  171. writer.Key("videoPath");
  172. writer.String(videoPath.c_str());
  173. writer.EndObject();
  174. }
  175. };