输出结果优化

Hu Chunming
1 parent 502f9d10
Showing 3 changed files with 75 additions and 20 deletions
src/ai_engine_module/VehiclePlate.cpp
src/ai_platform/header.h
src/demo/demo.cpp
@@ -88,9 +88,29 @@ vector&lt;VehiclePlateResult&gt; VehiclePlate::process(vector&lt;DeviceMemory*&gt; vec_gpuMe
         memcpy(res.recg, plate_info.recg, sizeof(PlateNum) * 8);
         res.num_score = plate_info.num_score;
         res.type = plate_info.type;
+
         res.state = plate_info.state;
         res.state_score = plate_info.state_score;
+        int space_char_num = 0;
+        for (size_t i = 0; i < 8; i++)
+        {
+            if (fabs(res.recg[i].maxprob) < 0.0001) {
+                space_char_num++ ;
+            }
+        }
+
+        if (res.state == 0)
+        {
+            if (space_char_num == 0) {
+                res.state = 1;
+                res.state_score = res.detect_score;
+            } else if (space_char_num > 0 && space_char_num < PLATENUM) {
+                res.state = 2;
+                res.state_score = res.detect_score;
+            }
+        }
+
         vec_res.push_back(res);
     }
@@ -46,6 +46,8 @@ typedef struct hp_cls_info         //分类结果结构体
     hp_cls_info() : res_index(0), res_prob(0) {};
 } hp_cls_info;
+
+static string vechicle_pendant[] = {"车头", "车尾", "车窗", "司机", "人脸", "安全带", "遮阳板", "年检标", "挂件", "纸巾盒", "转经筒", "打电话", "天窗", "行李架", "吸烟", "车前脸", "车后脸", "后视镜", "临时车牌", "车标", "三角标志", "三角危险", "长方形危险"};
 typedef struct VPDInfo
 {
 	int left_;
@@ -64,6 +66,8 @@ struct PlateNum
     int index;
 };
+static string plate_type[] = {"单排蓝色 ", "单排黄色 ", "单排白色 ", "单排黑色", " 双排黄色 ", "双排白色 ", "新能源黄绿色 ", "新能源白绿色"};
+static string plate_state[] = {"无车牌", " 正常车牌", " 遮挡车牌"};
 struct VehiclePlateResult {
     int left_;
 	int top_;
@@ -79,6 +83,8 @@ struct VehiclePlateResult {
     float state_score;//车牌状态置信度
 };
+// 0-black 1-blue 2-brown 3-green 4-grey 5-orange 6-pink 7-purple 8-red 9-silver 10-white 11-yellow
+static string vehicle_color[] = {"black", "blue", "brown", "green", "grey", "orange", "pink", "purple", "red", "silver", "white", "yellow"};
 typedef struct ObjectData {
     string task_id;              //该物体属于的任务ID号
     int task_frame_count;     //该物体当前出现的帧号
@@ -6,6 +6,8 @@
 #include <stdlib.h>
 #include <unistd.h>
 #include <iostream>
+#include <fstream>
+#include <string>
 #include "opencv2/opencv.hpp"
 #include "opencv2/imgcodecs/legacy/constants_c.h"
@@ -164,8 +166,13 @@ string createTask(void *handle, int gi, bool bFlag = true){
 void show_result(ObjectData obj)
 {
-  printf("%s  %d \n", obj.task_id.c_str(), obj.task_frame_count);
-  printf("%d %d (%d,%d,%d,%d) %lf \n", obj.index, obj.object_id, obj.left, obj.top, obj.right, obj.bottom, obj.confidence);
+  string str_dir = "./res/result/";
+  string file_name = obj.task_id + "_" + to_string(obj.object_id) + "_" + obj.str_ts_ms;
+  string data_file_name =  str_dir + file_name + ".txt";
+  std::ofstream res_file(data_file_name);
+
+  res_file << obj.task_id << "  " << obj.task_frame_count << endl;
+  res_file << obj.index << " " << obj.object_id << " " << obj.left << " " << obj.top << " " << obj.right << " " << obj.bottom << " " << obj.confidence << endl;
   cv::Mat image = cv::imread(obj.ori_pic_path);
@@ -175,26 +182,46 @@ void show_result(ObjectData obj)
   if (obj.index == 0)
   {// person
-    for (size_t i = 0; i < HP_ATTRI_SIZE; i++)
-    {
-      auto hp_cls = obj.hp_cls[i];
-      printf("(%d, %f)  ", hp_cls.res_index, hp_cls.res_prob);
-    }
-    printf("\n");
+    res_file << "person:" << endl;
+    res_file << "hair: " << hair[obj.hp_cls[0].res_index] << " " << obj.hp_cls[0].res_prob << endl;
+    res_file << "hair_color: " << hair_color[obj.hp_cls[1].res_index] << " " << obj.hp_cls[1].res_prob << endl;
+    res_file << "eye_glass: " << eye_glass[obj.hp_cls[2].res_index] << " " << obj.hp_cls[2].res_prob << endl;
+    res_file << "mask: " << mask[obj.hp_cls[3].res_index] << " " << obj.hp_cls[3].res_prob << endl;
+    res_file << "up cloth: " << up_cloth[obj.hp_cls[4].res_index] << " " << obj.hp_cls[4].res_prob << endl;
+    res_file << "up_cloth_color: " << up_cloth_color[obj.hp_cls[5].res_index] << " " << obj.hp_cls[5].res_prob << endl;
+    res_file << "clothing context: " << clothing_text[obj.hp_cls[6].res_index] << " " << obj.hp_cls[6].res_prob << endl;
+    res_file << "down cloth: " << down_cloth[obj.hp_cls[7].res_index] << " " << obj.hp_cls[7].res_prob << endl;
+    res_file << "down cloth color: " << down_cloth_color[obj.hp_cls[8].res_index] << " " << obj.hp_cls[8].res_prob << endl;
+    res_file << "bag: " << bag[obj.hp_cls[9].res_index] << " " << obj.hp_cls[9].res_prob << endl;
+    res_file << "sex: " << sex[obj.hp_cls[10].res_index] << " " << obj.hp_cls[10].res_prob << endl;
+    res_file << "age: " << age[obj.hp_cls[11].res_index] << " " << obj.hp_cls[11].res_prob << endl;
+    res_file << "viewpoint: " << viewpoint[obj.hp_cls[12].res_index] << " " << obj.hp_cls[12].res_prob << endl;
+    res_file << "umbrella: " << umbrella[obj.hp_cls[13].res_index] << " " << obj.hp_cls[13].res_prob << endl;
+    res_file << "hold_baby: " << hold_baby[obj.hp_cls[14].res_index] << " " << obj.hp_cls[14].res_prob << endl;
+    res_file << "person state: " << personstate[obj.hp_cls[15].res_index] << " " << obj.hp_cls[15].res_prob << endl;
+
     for (size_t i = 0; i < HUMANREID_FEATURE_SIZE; i++)
     {
       auto hp_cls = obj.hp_feature[i];
-      printf("%f ", obj.hp_feature[i]);
+      res_file << obj.hp_feature[i] << " ";
     }
-    printf("\n");
+    res_file << endl;
   } 
   else if (obj.index >= 4 && obj.index <= 8)
   {// vehicle
-    printf("vehicle color:%d  prob:%f \n", obj.vehicle_color_index, obj.vehicle_color_prob);
+    res_file << "vehicle:" << endl;
+    if (obj.vehicle_color_index >= 0)
+    {
+      res_file << "color:" << vehicle_color[obj.vehicle_color_index] << " " << obj.vehicle_color_prob << endl;
+    }
+    
     for (size_t i = 0; i < obj.vec_vpd_cls.size(); i++)
     {
       auto info = obj.vec_vpd_cls[i];
-      printf("%d %f (%d, %d, %d, %d) \n", info.index, info.confidence, info.left_, info.top_, info.width_, info.height_);
+      if (info.index >= 0)
+      {
+        res_file << vechicle_pendant[info.index] << " " << info.confidence << " " << info.left_ << " " << info.top_ << " " << info.width_ << " " << info.height_ << endl;
+      }
       cv::Rect rc(info.left_, info.top_, info.width_, info.height_);
       const cv::Scalar color(0, 255, 255);
@@ -202,22 +229,24 @@ void show_result(ObjectData obj)
     }
     auto& plate_info = obj.vehicle_plate;
-    printf("plate : (%d, %d, %d, %d)  %f \n", plate_info.left_, plate_info.top_, plate_info.width_, plate_info.height_, plate_info.detect_score );
+    res_file << "plate:" << plate_info.left_ << " " << plate_info.top_ << " " << plate_info.width_ << " " << plate_info.height_ << " " << plate_info.detect_score << endl;
     cv::Rect rc(plate_info.left_, plate_info.top_, plate_info.width_, plate_info.height_);
     const cv::Scalar color(0, 255, 0);
     cv::rectangle(image, rc, color);
+    res_file << "plate number:" ;
     for (size_t i = 0; i < 8; i++)
     {
-      printf("%s:%f  ", plate_info.recg[i].character, plate_info.recg[i].maxprob);
+      res_file << plate_info.recg[i].character << " " << plate_info.recg[i].maxprob << " ";
     }
-    printf("\n");
-    printf("num_score: %f \n", plate_info.num_score);
-    printf("type: %d \n", plate_info.type);
-    printf("state:%d  %f\n", plate_info.state, plate_info.state_score);
+    res_file << endl;
+    res_file << "num_score:" << plate_info.num_score << endl;
+    res_file << "type:" << plate_type[plate_info.type] << endl;
+    res_file << "state:" << plate_state[plate_info.state] << " " << plate_info.state_score << endl;
   }
-  printf("\n");
+
+  res_file.close();
-  string result_path = "./res/result/" +  obj.task_id + "_" + to_string(obj.object_id) + "_" + obj.str_ts_ms + ".jpg";
+  string result_path = str_dir + file_name + ".jpg";
   cv::imwrite(result_path, image);
   image.release();
 }