Procházet zdrojové kódy

version 3.4.5 msg: 加入InfineFilter模块, 负责过滤制定物品、过滤夜间发光物体提高夜间识别率。优化OSD识别框样式

Your Name před 3 roky
rodič
revize
348d695fe5

+ 5 - 0
modules/CMakeLists.txt

@@ -48,6 +48,7 @@ option(build_Base64          "build module build_Base64" ON)
 option(build_recorder        "build module build_recorder" ON)
 option(build_udp             "build module build_udp" ON)
 option(build_locus           "build_locus" ON)
+option(build_InfineFilter    "build_InfineFilter" ON)
 option(WITH_RTSP             "with rtsp" ON)
 option(WITH_FFMPEG           "with ffmpeg" ON)
 option(WITH_FFMPEG_AVDEVICE  "with ffmpeg avdevice" OFF)
@@ -257,6 +258,10 @@ if(build_udp)
   list(APPEND module_list udp)
   install(DIRECTORY udp/include/ DESTINATION include)
 endif()
+if(build_InfineFilter)
+  list(APPEND module_list InfineFilter)
+  install(DIRECTORY InfineFilter/include/ DESTINATION include)
+endif()
   
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})

+ 87 - 0
modules/InfineFilter/include/InfineFilter.h

@@ -0,0 +1,87 @@
+/*
+ * @Description: 
+ * @Version: 1.0
+ * @Autor: lishengyin
+ * @Date: 2022-03-22 14:00:35
+ * @LastEditors: lishengyin
+ * @LastEditTime: 2022-03-25 09:58:35
+ */
+#ifndef __INFINEFILTER_H_
+#define __INFINEFILTER_H_
+
+
+#include <memory>
+#include <string>
+#include <vector>
+#include <map>
+#include <list>
+#include <queue>
+
+#include "cnstream_frame.hpp"
+#include "cnstream_module.hpp"
+#include "easyinfer/model_loader.h"
+#include "easytrack/easy_track.h"
+#include "cnstream_frame_va.hpp"
+
+#include "opencv2/highgui/highgui.hpp"
+#include "opencv2/imgproc/imgproc.hpp"
+#include "opencv2/dnn.hpp"
+#include "opencv2/opencv.hpp"
+#include "opencv2/highgui.hpp"
+#include "opencv2/imgcodecs/imgcodecs.hpp"
+
+#include <rapidjson/document.h>
+#include <rapidjson/rapidjson.h>
+#include <rapidjson/stringbuffer.h>
+#include <rapidjson/writer.h>
+#include <ctime>
+#include <iostream>
+#include <time.h>
+#include <cstdlib>
+#include <dirent.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+
+using namespace std;
+
+
+namespace cnstream{
+    // Pointer for frame info
+    using CNFrameInfoPtr = std::shared_ptr<CNFrameInfo>;
+    /// Pointer for infer object
+    using CNInferObjectPtr = std::shared_ptr<CNInferObject>;
+
+    class InfineFilter : public cnstream::Module, public cnstream::ModuleCreator<InfineFilter> {
+        using super = cnstream::Module;
+        
+    public:
+        explicit InfineFilter(const std::string &name) : super(name) {}
+        bool Open(cnstream::ModuleParamSet paramSet) override;
+
+        void Close() override { std::cout << this->GetName() << " Close called" << std::endl; }
+
+        int Process(std::shared_ptr<cnstream::CNFrameInfo> data) override;
+
+        bool filter(cv::Mat image, std::shared_ptr<cnstream::CNInferObject>& object);
+        
+        bool getNight(cv::Mat image);
+
+        std::string getTime();
+
+    private:
+        InfineFilter(const InfineFilter &) = delete;
+        InfineFilter &operator=(InfineFilter const &) = delete;
+        
+        vector<std::string> labels;
+
+        bool saveFilterResult = true;
+        std::string files = "./output/";
+
+        int start = 21;
+        int end = 6;
+    };
+};
+
+
+
+#endif

+ 189 - 0
modules/InfineFilter/src/InfineFilter.cpp

@@ -0,0 +1,189 @@
+/*
+ * @Description: 
+ * @Version: 1.0
+ * @Autor: lishengyin
+ * @Date: 2022-03-22 14:00:46
+ * @LastEditors: lishengyin
+ * @LastEditTime: 2022-03-25 09:56:36
+ */
+#include "InfineFilter.h"
+
+namespace cnstream{
+
+    bool InfineFilter::Open(cnstream::ModuleParamSet paramSet){
+        if (paramSet.find("FilterId") != paramSet.end()) {
+            std::string json = paramSet["FilterId"];
+            rapidjson::Document doc;
+            if (doc.Parse<rapidjson::kParseCommentsFlag>(json.c_str()).HasParseError()) {
+                return false;
+            }
+            const auto end = doc.MemberEnd();
+            if(end == doc.FindMember("label")|| !(doc["label"].IsArray())){
+                return false;
+            }else{
+                const rapidjson::Value& objs = doc["label"];
+                for(size_t i = 0; i < objs.Size(); i++){
+                    labels.push_back(objs[i].GetString());
+                }
+            }
+            for(auto iter = labels.begin(); iter != labels.end(); iter++){
+                LOGE(InfineFilter) << *iter << ",";
+            }
+        }
+        if(paramSet.find("files") != paramSet.end()){
+            this->files = paramSet["files"];
+        }
+        if(paramSet.find("saveFilterResult") != paramSet.end()){
+            this->saveFilterResult = paramSet["saveFilterResult"] == "true" ? true : false;
+        }
+        if(paramSet.find("lock_period") != paramSet.end()){
+            std::string json = paramSet["lock_period"];
+            rapidjson::Document doc;
+            if (doc.Parse<rapidjson::kParseCommentsFlag>(json.c_str()).HasParseError()) {
+                return false;
+            }
+            const auto end = doc.MemberEnd();
+            if(end == doc.FindMember("start")|| !(doc["start"].IsInt())){
+                return false;
+            }else{
+                this->start = doc["start"].GetInt();
+            }
+            if(end == doc.FindMember("end")|| !(doc["end"].IsInt())){
+                return false;
+            }else{
+                this->end = doc["end"].GetInt();
+            }
+        }
+        return true;
+    }
+
+    int InfineFilter::Process(std::shared_ptr<cnstream::CNFrameInfo> data){
+        CNDataFramePtr frame = data->collection.Get<CNDataFramePtr>(kCNDataFrameTag);
+        if (frame->width < 0 || frame->height < 0) {
+            LOGE(InfineFilter) << "InfineFilter module processed illegal frame: width or height may < 0.";
+            return -1;
+        }
+        
+        CNInferObjsPtr objs_holder = nullptr;
+        if (data->collection.HasValue(kCNInferObjsTag)) {
+            objs_holder = data->collection.Get<CNInferObjsPtr>(kCNInferObjsTag);
+        }
+        
+        if(objs_holder->objs_.size() == 0){
+            this->TransmitData(data);
+            return 0;
+        }
+
+        // 过滤某些Id
+        for(auto iter = objs_holder->objs_.begin(); iter != objs_holder->objs_.end(); ){
+            std::shared_ptr<cnstream::CNInferObject> object = *iter;
+            if (!object) continue;
+            auto it = find(this->labels.begin(), this->labels.end(), object->id);
+            if(it != this->labels.end()){
+                iter = objs_holder->objs_.erase(iter);
+            }else{
+                iter++;
+            }
+        }
+        
+        // 过滤星星、月亮、飞机
+        for(auto iter = objs_holder->objs_.begin(); iter != objs_holder->objs_.end();){
+            std::shared_ptr<cnstream::CNInferObject> object = *iter;
+            if (!object) continue;
+            if(!this->filter(frame->ImageBGR(), object)){
+                iter = objs_holder->objs_.erase(iter);
+            }else{
+                iter++;
+            }
+        }   
+        return 0;
+    }
+
+    /**
+     * @description: 过滤器
+     * @param {Mat} image
+     * @return {*}
+     */    
+    bool InfineFilter::filter(cv::Mat image, std::shared_ptr<cnstream::CNInferObject>& object){
+        if(!getNight(image)) return true;
+        
+        cv::Mat img = image.clone();
+        // 腐蚀
+        erode(img, img, cv::Mat());
+
+        cv::Rect rect = cv::Rect(object->bbox.x * image.cols,object->bbox.y * image.rows, object->bbox.w * image.cols, object->bbox.h * image.rows);
+        cv::Mat dst = img(rect);
+        cv::Mat imgThresholded;
+        // 转换为二值图
+        cv::inRange(dst, cv::Scalar(128, 128, 128), cv::Scalar(255, 255, 255), imgThresholded);
+        cv::threshold(imgThresholded, imgThresholded, 1, 255, cv::THRESH_BINARY);
+        
+        // 寻找轮廓
+        vector<vector<cv::Point>> contours;
+        vector<cv::Vec4i> hierarchy;
+        cv::findContours(imgThresholded, contours, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_SIMPLE);
+
+        if(contours.empty()) return true;
+        // 获取最大轮廓
+        double max_area = cv::contourArea(cv::InputArray(contours[0]), false);
+        for(unsigned long long i = 1; i < contours.size(); i++){
+            double temp_area = cv::contourArea(cv::InputArray(contours[i]), false);
+            if(max_area < temp_area){
+                max_area = temp_area;
+            }
+        }
+        double b = max_area / (object->bbox.w * image.cols * object->bbox.h * image.rows);
+        // 计算轮廓比例
+        if(b > 0.5){
+            if(saveFilterResult){
+                cv::Point top_left;
+                cv::Point bottom_right;
+                top_left.x = object->bbox.x * image.cols;
+                top_left.y = object->bbox.y * image.rows;
+                bottom_right.x = top_left.x + (object->bbox.w * image.cols);
+                bottom_right.y = top_left.y + (object->bbox.h * image.rows);
+                cv::Point logo_pos(5, image.rows - 5);
+                cv::Scalar color(200, 200, 200);
+                cv::putText(img, std::to_string(b), logo_pos, 0, 1, color, 2);
+                cv::rectangle(img, top_left, bottom_right, cv::Scalar(0, 255, 0), 4);
+                if(opendir(this->files.c_str()) == NULL){
+                    mkdir((this->files).c_str(),S_IRWXU|S_IRWXG|S_IRWXO);
+                } 
+                cv::imwrite(this->files + getTime() + ".jpg", img);
+            }
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * @description: 判断时间是否为晚上
+     * @param {*}
+     * @return {*}
+     */    
+    bool InfineFilter::getNight(cv::Mat image){
+        time_t now = time(0);
+        tm *gmtm = gmtime(&now);
+        if(gmtm->tm_hour > this->start ||  gmtm->tm_hour < this->end){
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * @description: 获取时间
+     * @param {*}
+     * @return {*}
+     */    
+    std::string InfineFilter::getTime(){
+        std::string nowtime;
+        time_t rawtime;
+        char ctime[80];
+        struct tm *info;
+        time(&rawtime);
+        info = localtime(&rawtime);
+        strftime(ctime, 80, "%Y-%m-%d_%H:%M:%S", info);
+        nowtime = ctime;
+        return nowtime;
+    }
+};

+ 0 - 1
modules/locus/include/locus.h

@@ -7,7 +7,6 @@
 #include <list>
 #include <queue>
 
-
 #include "cnstream_frame.hpp"
 #include "cnstream_module.hpp"
 #include "easyinfer/model_loader.h"

+ 29 - 2
modules/osd/src/cnosd.cpp

@@ -119,6 +119,8 @@ void CnOsd::DrawLabel(cv::Mat image, const CNInferObjsPtr& objs_holder, std::vec
     return;
   }
 
+  bool drawText = true;
+  if(objs_holder->objs_.size() > 5) drawText = false;
   for (uint32_t i = 0; i < objs_holder->objs_.size(); ++i) {
     std::shared_ptr<cnstream::CNInferObject> object = objs_holder->objs_[i];
     if (!object) continue;
@@ -155,7 +157,7 @@ void CnOsd::DrawLabel(cv::Mat image, const CNInferObjsPtr& objs_holder, std::vec
     } else {
       LOGD(OSD) << "Draw Label and Score: " << text;
     }
-    DrawText(image, bottom_left, text, color);
+    if(drawText) DrawText(image, bottom_left, text, color);
 
     // draw secondary inference information
     int label_bottom_y = 0;
@@ -204,7 +206,32 @@ int CnOsd::GetLabelId(const std::string& label_id_str) const {
 
 void CnOsd::DrawBox(cv::Mat image, const cv::Point& top_left, const cv::Point& bottom_right,
                     const cv::Scalar& color) const {
-  cv::rectangle(image, top_left, bottom_right, color, CalcThickness(image.cols, box_thickness_));
+  float w =bottom_right.x - top_left.x;
+  float h =bottom_right.y - top_left.y;
+  int wscale = w * .25f;
+  int hscale = h * .25f;
+  cv::Point top_right(bottom_right.x,top_left.y);
+  cv::Point bottom_left(top_left.x,bottom_right.y);
+  // -- 
+  cv::line(image,top_left,cv::Point(top_left.x + wscale,top_left.y),color, CalcThickness(image.cols, box_thickness_));
+  // |
+  cv::line(image,top_left,cv::Point(top_left.x ,top_left.y + hscale),color, CalcThickness(image.cols, box_thickness_));
+  //      --
+  cv::line(image,cv::Point(top_right.x - wscale,top_right.y),top_right,color, CalcThickness(image.cols, box_thickness_));
+  //        |
+  cv::line(image,top_right,cv::Point(top_right.x,top_right.y + hscale),color, CalcThickness(image.cols, box_thickness_));
+  //
+  // __
+  cv::line(image,bottom_left,cv::Point(bottom_left.x+wscale,bottom_left.y),color, CalcThickness(image.cols, box_thickness_));
+  //
+  //|
+  cv::line(image,bottom_left,cv::Point(bottom_left.x,bottom_left.y - hscale),color, CalcThickness(image.cols, box_thickness_));
+  //
+  //        |
+  cv::line(image,bottom_right,cv::Point(bottom_right.x,bottom_right.y - hscale),color, CalcThickness(image.cols, box_thickness_));
+  
+  cv::line(image,bottom_right,cv::Point(bottom_right.x - wscale,bottom_right.y),color, CalcThickness(image.cols, box_thickness_));
+  // cv::rectangle(image, top_left, bottom_right, color, CalcThickness(image.cols, box_thickness_));
 }
 
 void CnOsd::DrawText(cv::Mat image, const cv::Point& bottom_left, const std::string& text, const cv::Scalar& color,

+ 18 - 0
source/cns_launcher/configs/InfineFilter.json

@@ -0,0 +1,18 @@
+{
+    "InfineFilter" : {
+        "class_name" : "cnstream::InfineFilter",
+        "parallelism" : 1,
+        "max_input_queue_size" : 20,
+        "custom_params" : {
+            "FilterId" : {
+                "label":[]
+            },
+            "files": "./output/",
+            "saveFilterResult" : "false",
+            "lock_period":{
+                "start": 21,
+                "end": 6
+            }
+        }
+    }
+}

+ 12 - 6
source/cns_launcher/object_tracking/detection_test.json

@@ -11,8 +11,14 @@
   
     "subgraph:object_detection" : {
       "config_path" : "../configs/yolov5_object_detection_test.json",
+      "next_modules" : ["subgraph:InfineFilter"]
+    },
+
+    "subgraph:InfineFilter" : {
+      "config_path" : "../configs/InfineFilter.json",
       "next_modules" : ["subgraph:object_tracking"]
     },
+
   
     "subgraph:object_tracking" : {
       "config_path" : "../configs/object_tracking_mlu220.json",
@@ -21,11 +27,6 @@
   
     "subgraph:udp" : {
       "config_path" : "../configs/udp.json",
-      "next_modules" : ["subgraph:kafka"]
-    },
-  
-    "subgraph:kafka" : {
-      "config_path" : "../configs/ipc.json",
       "next_modules" : ["subgraph:osd_label_map_coco"]
     },
   
@@ -41,9 +42,14 @@
   
     "subgraph:object_recorder" : {
       "config_path" : "../configs/recorder.json",
+      "next_modules" : ["subgraph:kafka"]
+    },
+    
+    "subgraph:kafka" : {
+      "config_path" : "../configs/ipc.json",
       "next_modules" : ["subgraph:sinker"]
     },
-  
+
     "subgraph:sinker" : {
       "config_path" : "../configs/sinker_configs/rtsp.json"
     }

binární
source/cns_launcher/object_tracking/output/output_stream_0.mp4


+ 1 - 1
source/common/postprocess/video_postprocess_yolov5.cpp

@@ -121,7 +121,7 @@ bool VideoPostprocYolov5::Execute(infer_server::InferData* output_data, const in
     obj->bbox.w = std::min(1.0f - obj->bbox.x, right - left);
     obj->bbox.h = std::min(1.0f - obj->bbox.y, bottom - top);
 
-    if (obj->bbox.h <= 0 || obj->bbox.w <= 0 || (obj->score < threshold_ && threshold_ > 0) || obj->id != "0") continue;
+    if (obj->bbox.h <= 0 || obj->bbox.w <= 0 || (obj->score < threshold_ && threshold_ > 0)) continue;
     std::lock_guard<std::mutex> objs_mutex(objs_holder->mutex_);
     objs.push_back(obj);
   }