完善pase_road

Hu Chunming
1 parent 01b357d4
Showing 2 changed files with 112 additions and 56 deletions
src/ai_engine_module/RoadSegAnalysis.cpp
src/ai_engine_module/RoadSegAnalysis.h
@@ -63,9 +63,9 @@ std::vector&lt;RoadInfo&gt; RoadSegAnalysis::detect(vector&lt;sy_img&gt; vec_img){
         for (int b = 0; b < batchsize; b++) {
             auto one_result = results[b];
-            RoadInfo one_road;
-            one_road.vec_line = parse_line(one_result, vec_img[b]);
-            one_road.vec_road = parse_road(one_result, vec_img[b]);
+            RoadInfo one_road = parse_road(one_result, vec_img[b]);
+            // one_road.vec_line = parse_line(one_result, vec_img[b]);
+            // one_road.vec_road = parse_road(one_road, one_result, vec_img[b]);
             one_road.vec_direct = parse_direct(one_result, vec_img[b]);
             vec_road_info.push_back(one_road);
         }
@@ -147,14 +147,14 @@ cv::Mat mask_to_rgb(cv::Mat img, cv::Mat mask) {
     return masks;
 }
-std::vector<SegInfo> RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
+RoadInfo RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
 {
     int src_width = src.w_;
     int src_height = src.h_;
     float scale_w = src_width / 640.0;
 	float scale_h = src_height / 360.0;
-    std::vector<SegInfo> vec_road;
+    RoadInfo one_road;
     std::vector<std::pair<std::vector<cv::Point>, int>> combined;
     lanes_process(one_result.reg_array, one_result.lane_count, combined);
@@ -167,88 +167,82 @@ std::vector&lt;SegInfo&gt; RoadSegAnalysis::parse_road(rs_result one_result, sy_img sr
     cv::Mat seg_output = seg_post_process(large_resolution, one_result.seg_array, combined, poly_masks, region_classes, lanes, cats, x_sort); //m_masks:mask前的结果  poly_masks后的结果               
     // cv::Mat image_lane(cv::Size(src_width, src_height), CV_8UC3, cv::Scalar(0, 0, 0));
-    // for (int i = 0; i < lanes.size(); ++i) {
-    //     LineInfo info;
-    //     const cv::Scalar color(seg_colors[i][0], seg_colors[i][1], seg_colors[i][2]);
-    //     for (size_t j = 0; j < lanes[i].size(); j++) {
-    //         cv::Point pt;
-    //         pt.x = lanes[i][j].x * scale_w;
-    //         pt.x = lanes[i][j].y * scale_h;
-    //         clip_xy(pt.x, pt.y, src_width, src_height);
-    //         info.vec_pt.push_back(pt);
-
-    //         if (j > 0) {
-    //             cv::line(image_lane, info.vec_pt[j-1], info.vec_pt[j], color, 2, 8);
-    //         }
+    for (int i = 0; i < lanes.size(); ++i) {
+        LineInfo info;
+        const cv::Scalar color(seg_colors[i][0], seg_colors[i][1], seg_colors[i][2]);
+        for (size_t j = 0; j < lanes[i].size(); j++) {
+            cv::Point pt;
+            pt.x = lanes[i][j].x * scale_w;
+            pt.y = lanes[i][j].y * scale_h;
+            clip_xy(pt.x, pt.y, src_width, src_height);
+            info.vec_pt.push_back(pt);
+
+            // if (j > 0) {
+            //     cv::line(image_lane, info.vec_pt[j-1], info.vec_pt[j], color, 2, 8);
+            // }
-    //     }
+        }
-    //     info.line_type = cats[i];
-    //     one_road.vec_line.push_back(info);
-    // }
-    cv::Mat cvImg(cv::Size(src_width, src_height), CV_8UC3, cv::Scalar(0, 0, 0));
-    cv::Mat tmp;
-    cv::resize(cvImg, tmp, cv::Size(640,360), 0, 0, cv::INTER_LINEAR);
-    cv::Mat vis_image = mask_to_rgb(tmp, seg_output);
-    cv::imwrite("./image_lane.jpg", vis_image);
-    seg_output.release();
-    cvImg.release();
+        info.line_type = cats[i];
+        one_road.vec_line.push_back(info);
+    }
+
+    // cv::imwrite("./image_lane.jpg", image_lane);
+    // seg_output.release();
+    // image_lane.release();
-    cv::Mat image(cv::Size(src_width, src_height), CV_8UC3, cv::Scalar(0, 0, 0));
+    // cv::Mat image(cv::Size(src_width, src_height), CV_8UC3, cv::Scalar(0, 0, 0));
     for (int i = 0; i < poly_masks.size(); ++i) {
         SegInfo seg_info;
         for (size_t j = 0; j < poly_masks[i].size(); j++) {
             cv::Point pt;
             pt.x = poly_masks[i][j].x * scale_w;
-            pt.x = poly_masks[i][j].y * scale_h;
+            pt.y = poly_masks[i][j].y * scale_h;
             clip_xy(pt.x, pt.y, src_width, src_height);
             seg_info.vec_pt.push_back(pt);
         }
         seg_info.seg_type = region_classes[i];
-        vec_road.push_back(seg_info);
+        one_road.vec_road.push_back(seg_info);
-        int k = seg_info.seg_type;
-        const cv::Scalar color(seg_colors[k][0], seg_colors[k][1], seg_colors[k][2]);
-        polylines(image, seg_info.vec_pt, true, color, 3, cv::LINE_AA);
+        // int k = seg_info.seg_type;
+        // const cv::Scalar color(seg_colors[k][0], seg_colors[k][1], seg_colors[k][2]);
+        // polylines(image, seg_info.vec_pt, true, color, 3, cv::LINE_AA);
     }
-    cv::imwrite("./image_road.jpg", image);
-    image.release();
+    // cv::imwrite("./image_road.jpg", image);
+    // image.release();
-    return vec_road;
+    return one_road;
 }
 std::vector<SegInfo> RoadSegAnalysis::parse_direct(rs_result one_result, sy_img src) {
-    std::vector<std::pair<std::vector<cv::Point>, int>> combined;
-    lanes_process(one_result.reg_array, one_result.lane_count, combined);
-    std::vector<std::vector<cv::Point>> poly_masks, lanes;
-    std::vector<int> region_classes, cats;
-    std::map<double, int> x_sort;
+
+    std::vector<std::vector<cv::Point>> direct_masks;
+    std::vector<int> direct_classes;
     bool large_resolution = false;
     if (src.w_ > 1920) large_resolution = true;
-
-    cv::Mat seg_output = seg_post_process(large_resolution, one_result.direct_seg, combined, poly_masks, region_classes, lanes, cats, x_sort); //m_masks:mask前的结果  poly_masks后的结果               
+    cv::Mat direct_output = direct_post_process(large_resolution, one_result.direct_seg, direct_masks, direct_classes);          
-    cv::Mat image(cv::Size(640, 360), CV_8UC3, cv::Scalar(0, 0, 0));
+    // cv::Mat image(cv::Size(640, 360), CV_8UC3, cv::Scalar(0, 0, 0));
     std::vector<SegInfo> vec_road;
-    for (int i = 0; i < poly_masks.size(); ++i) {
+    for (int i = 0; i < direct_masks.size(); ++i) {
         SegInfo seg_info;
-        seg_info.seg_type = region_classes[i];
-        seg_info.vec_pt = poly_masks[i];
+        seg_info.seg_type = direct_classes[i];
+        seg_info.vec_pt = direct_masks[i];
         vec_road.push_back(seg_info);
-        int k = seg_info.seg_type;
-        const cv::Scalar color(seg_colors[k][0], seg_colors[k][1], seg_colors[k][2]);
-        polylines(image, seg_info.vec_pt, true, color, 3, cv::LINE_AA);
+        // int k = seg_info.seg_type;
+        // const cv::Scalar color(seg_colors[k][0], seg_colors[k][1], seg_colors[k][2]);
+        // polylines(image, seg_info.vec_pt, true, color, 3, cv::LINE_AA);
     }
-    cv::imwrite("./image_direct.jpg", image);
-    image.release();
+    // cv::imwrite("./image_direct.jpg", image);
+    // image.release();
     return vec_road;
 }
@@ -679,4 +673,64 @@ cv::Mat RoadSegAnalysis::seg_post_process(bool large_resolution, unsigned char *
     }
 #endif
     return mask_rmlane;
-}
 \ No newline at end of file
+}
+
+cv::Mat RoadSegAnalysis::direct_post_process(bool large_resolution, unsigned char *direct_array, std::vector<std::vector<cv::Point>> &poly_masks, std::vector<int> &region_classes) {
+    std::vector<int> pred_cls; 
+    int h = 360, w = 640, step_size = h*w, seg_min_region_area = 512;
+    cv::Mat mask_rmlane = cv::Mat_<int>(h,w); //车道线区域置为背景
+	for (int i = 0; i < h; ++i) {
+		for (int j = 0; j < w; ++j) {
+			int max_cls = direct_array[(i * w + j)];
+            pred_cls.push_back(max_cls);
+            mask_rmlane.ptr<int>(i)[j] = max_cls;
+		}
+	}
+
+#if 1 /*mask远处区域*/
+    cv::Mat mask_black = mask_rmlane.clone();
+    if (large_resolution)   mask_black(cv::Rect(0, 0, w, int(h * 0.14))) = 0;  
+    else    mask_black(cv::Rect(0, 0, w, int(h * 0.22))) = 0;           
+    mask_black(cv::Rect(0, int(h * 0.95), w, int(h * 0.05))) = 0;
+    mask_black(cv::Rect(0, 0, int(w * 0.02), h)) = 0;
+    mask_black(cv::Rect(int(w * 0.95), 0, int(w * 0.05), h)) = 0;
+    mask_rmlane = mask_black;
+#endif
+    //去重获取预测到的类别
+    std::vector<int> labels(pred_cls);
+    std::sort(labels.begin(),labels.end());
+    labels.erase(std::unique(labels.begin(),labels.end()),labels.end());
+ 
+    //求道路区域
+    for(auto cat: labels) {
+        cv::Mat b_masks, measure_labels, stats, centroids;
+        cv::Mat n_masks = cv::Mat_<int>(h,w);   n_masks = cv::Scalar(255);
+        cv::compare(mask_rmlane, cat, b_masks, cv::CMP_EQ); //将mask_rmlane元素逐个和cat比较，相同255，不同0
+        //连通域计算
+        int nccomps = cv::connectedComponentsWithStats (b_masks, measure_labels, stats, centroids, 8);
+        for(int i = 1; i < nccomps; i++ ) {	//跳过背景信息0
+            if (stats.at<int>(i, cv::CC_STAT_AREA) < seg_min_region_area) { //移除过小的区域，并将对应位置置为0
+                cv::Mat comparison_result;
+                cv::compare(measure_labels, cv::Scalar(i), comparison_result, cv::CMP_EQ); //相等为255不等为0
+                n_masks.setTo(0, comparison_result); // 将comparison_result中非零区域置为0
+                cv::bitwise_and(mask_rmlane,n_masks,mask_rmlane);
+                continue;
+            }
+
+            cv::Mat region_mask;
+            cv::bitwise_and(measure_labels,n_masks,measure_labels);/////
+            cv::compare(measure_labels, i, region_mask, cv::CMP_EQ);
+
+            std::vector<std::vector<cv::Point>> contours;
+            cv::findContours(region_mask,contours,cv::noArray(),cv::RETR_TREE,cv::CHAIN_APPROX_SIMPLE);
+            std::vector<cv::Point> contours_poly;
+            for (int j = 0; j < contours.size();j++) {
+                cv::approxPolyDP(cv::Mat(contours[j]), contours_poly, 10, true);
+                if (contours_poly.size() <= 2) continue;
+                poly_masks.push_back(contours_poly);
+                region_classes.push_back(cat);  
+            }
+        }    
+    }
+    return mask_rmlane;
+}
@@ -33,7 +33,7 @@ private:
     std::vector<LineInfo> parse_line(rs_result one_result, sy_img src);
-    std::vector<SegInfo> parse_road(rs_result one_result, sy_img src);
+    RoadInfo parse_road(rs_result one_result, sy_img src);
     std::vector<SegInfo> parse_direct(rs_result one_result, sy_img src);
@@ -47,6 +47,8 @@ private:
     int Mask2LanePoints(const cv::Mat& pred, std::vector<std::vector<cv::Point>>&lanes, std::vector<int>& cats);
     cv::Mat seg_post_process(bool large_resolution, unsigned char *seg_array, std::vector<std::pair<std::vector<cv::Point>, int>> combined, std::vector<std::vector<cv::Point>> &poly_masks, std::vector<int> &region_classes, std::vector<std::vector<cv::Point>> &lanes, std::vector<int> &cats, std::map<double, int> &x_sort);
+    cv::Mat direct_post_process(bool large_resolution, unsigned char *direct_array, std::vector<std::vector<cv::Point>> &poly_masks, std::vector<int> &region_classes);
+
 private:
     void* m_handle{nullptr};
     aclrtContext ctx{nullptr};