修改 parse_line 实现

Hu Chunming
1 parent 5fc0b809
Showing 2 changed files with 121 additions and 23 deletions
src/ai_engine_module/RoadSegAnalysis.cpp
src/ai_engine_module/RoadSegAnalysis.h
@@ -35,7 +35,7 @@ int RoadSegAnalysis::init(int devId, std::string sdk_root){
 std::vector<RoadInfo> RoadSegAnalysis::detect(vector<sy_img> vec_img){
-    std::vector<RoadInfo> vec_road;
+    std::vector<RoadInfo> vec_road_info;
     const int batchsize = vec_img.size();
     rs_result results[batchsize];
@@ -63,9 +63,11 @@ 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 = parse_road(one_result, vec_img[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]);
             one_road.vec_direct = parse_direct(one_result, vec_img[b]);
-            vec_road.push_back(one_road);
+            vec_road_info.push_back(one_road);
         }
     } while (0);
@@ -77,7 +79,7 @@ std::vector&lt;RoadInfo&gt; RoadSegAnalysis::detect(vector&lt;sy_img&gt; vec_img){
         results[b].direct_seg = NULL;
     }
-    return vec_road;
+    return vec_road_info;
 }
 void test(vector<int> vec_type, std::vector<cv::Point> vec_pt, std::string file_name){
@@ -93,9 +95,66 @@ void test(vector&lt;int&gt; vec_type, std::vector&lt;cv::Point&gt; vec_pt, std::string file_
     image.release();
 }
-RoadInfo RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
+void clip_xy(int& x, int& y, int width, int height){
+    if (x < 0) {
+        x = 0;
+    }
+
+    if (y < 0) {
+        y = 0;
+    }
+    
+    if (x >= width) {
+        x = width -1;
+    }
+    
+    if (y >= height) {
+        y = height -1;
+    }
+}
+
+std::vector<LineInfo> RoadSegAnalysis::parse_line(rs_result one_result, sy_img src) {
+	float scale_w = src.w_ / 640.0;
+	float scale_h = src.h_ / 360.0;
+    std::vector<std::pair<std::vector<cv::Point>, int>> combined;
+    lanes_process(one_result.reg_array, one_result.lane_count, combined, scale_w, scale_h);
+
+    std::vector<LineInfo> vec_line;
+    for (const auto& lane_info : combined) {
+        LineInfo info;
+        info.vec_pt = lane_info.first;
+        info.line_type = lane_info.second;
+        vec_line.push_back(info);
+    }
+    
+    return vec_line;
+}
+
+cv::Mat mask_to_rgb(cv::Mat img, cv::Mat mask) {
+    cv::Mat masks = img.clone();
+    int reg_cls = 9;
+    for (int i = 0; i < masks.rows; i++) {
+        for (int j = 0; j < masks.cols; j++) {
+            for (int k = 1; k < reg_cls; k++) {
+                if (mask.at<int>(i,j) == k) {
+                    masks.at<cv::Vec3b>(i,j)[0] = seg_colors[k][0];
+                    masks.at<cv::Vec3b>(i,j)[1] = seg_colors[k][1];
+                    masks.at<cv::Vec3b>(i,j)[2] = seg_colors[k][2];
+                }
+            }
+        }
+    }
+    return masks;
+}
+
+std::vector<SegInfo> RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
 {
-    RoadInfo one_road;
+    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;
     std::vector<std::pair<std::vector<cv::Point>, int>> combined;
     lanes_process(one_result.reg_array, one_result.lane_count, combined);
@@ -103,25 +162,62 @@ RoadInfo RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
     std::vector<int> region_classes, cats;
     std::map<double, int> x_sort;
     bool large_resolution = false;
-    if (src.w_ > 1920) large_resolution = true;
+    if (src_width > 1920) large_resolution = true;
     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后的结果               
-    for (int i = 0; i < lanes.size(); ++i) {
-        LineInfo info;
-        info.vec_pt = lanes[i];
-        info.line_type = cats[i];
-        one_road.vec_line.push_back(info);
-    } 
+    // 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);
+    //         }
+            
+    //     }
+        
+    //     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();
+    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;
-        seg_info.vec_pt = poly_masks[i];
+
+        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;
+            clip_xy(pt.x, pt.y, src_width, src_height);
+            seg_info.vec_pt.push_back(pt);
+        }
+
         seg_info.seg_type = region_classes[i];
-        one_road.vec_road.push_back(seg_info);
+        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);
     }
-    return one_road;
+    cv::imwrite("./image_road.jpg", image);
+    image.release();
+
+    return vec_road;
 }
 std::vector<SegInfo> RoadSegAnalysis::parse_direct(rs_result one_result, sy_img src) {
@@ -136,7 +232,7 @@ std::vector&lt;SegInfo&gt; RoadSegAnalysis::parse_direct(rs_result one_result, sy_img 
     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 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) {
@@ -146,13 +242,13 @@ std::vector&lt;SegInfo&gt; RoadSegAnalysis::parse_direct(rs_result one_result, sy_img 
         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("./direct.jpg", image);
-    // image.release();
+    cv::imwrite("./image_direct.jpg", image);
+    image.release();
     return vec_road;
 }
@@ -31,7 +31,9 @@ public:
 private:
     int release();
-    RoadInfo parse_road(rs_result one_result, sy_img src);
+    std::vector<LineInfo> parse_line(rs_result one_result, sy_img src);
+
+    std::vector<SegInfo> parse_road(rs_result one_result, sy_img src);
     std::vector<SegInfo> parse_direct(rs_result one_result, sy_img src);