src/ai_engine_module/RoadSegAnalysis.cpp

#include "RoadSegAnalysis.h"
uint8_t seg_colors[][3] = { {0, 0, 0}, {0, 255, 255}, {128, 255, 0}, {255, 128, 0}, {128, 0, 255}, {255, 0, 128}, {0, 128, 255}, {0, 255, 128}, {128, 255, 255}};
uint8_t lane_colors[][3] = { {0, 0, 0}, {255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255}, {0, 255, 255}, {128, 255, 0}, {255, 128, 0}};
RoadSegAnalysis::RoadSegAnalysis(/* args */)
{
}
RoadSegAnalysis::~RoadSegAnalysis()
{
    release();
}
int RoadSegAnalysis::init(int devId, std::string sdk_root){
    ACL_CALL(aclrtCreateContext(&ctx, devId), SY_SUCCESS, SY_FAILED);
    std::string model_path = sdk_root + "/models/road_seg/tzroad_seg240108_310p.om";
    rs_param param;
    param.modelNames = (char*)model_path.data();
    param.thresld = 0.25;
    param.devId = devId;
    LOG_INFO("rs_init start");
    int ret = rs_init(&m_handle, param);
    if (ret != 0) {
        return -1;
	}
    LOG_INFO("rs_init success");
    return SY_SUCCESS;
}
std::vector<RoadInfo> RoadSegAnalysis::detect(vector<sy_img> vec_img){
    std::vector<RoadInfo> vec_road;
    const int batchsize = vec_img.size();
    rs_result results[batchsize];
    const int fea_size = 360*640;
    for (int b = 0; b < batchsize; b++) {
        results[b].seg_array = new unsigned char[fea_size];
        results[b].direct_seg = new unsigned char[fea_size];
    }
    int ret = SY_FAILED;
    do
    {
        ret = aclrtSetCurrentContext(ctx);
        if (SY_SUCCESS != ret) {
            printf("aclrtSetCurrentContext failed!");
            break;
        }
        ret = rs_batch(m_handle, vec_img.data(), batchsize, results);
        if (SY_SUCCESS != ret) {
            printf("rs_batch failed!");
            break;
        }
        
        for (int b = 0; b < batchsize; b++) {
            auto one_result = results[b];
            RoadInfo one_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);
        }
    } while (0);
     for (int b = 0; b < batchsize; b++) {
        delete[] results[b].seg_array; 
        results[b].seg_array = NULL;
        delete[] results[b].direct_seg; 
        results[b].direct_seg = NULL;
    }
    return vec_road;
}
// void test(){
//     cv::Mat image(cv::Size(640, 360), CV_8UC3, cv::Scalar(0, 0, 0));
//     for (int i = 0; i < poly_masks.size(); ++i) {
//         SegInfo seg_info;
//         seg_info.seg_type = region_classes[i];
//         seg_info.vec_pt = poly_masks[i];
//         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("./road.jpg", image);
//     image.release();
// }
RoadInfo RoadSegAnalysis::parse_road(rs_result one_result, sy_img src)
{
    RoadInfo one_road;
    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;
    bool large_resolution = false;
    if (src.w_ > 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);
    } 
    
    for (int i = 0; i < poly_masks.size(); ++i) {
        SegInfo seg_info;
        seg_info.vec_pt = poly_masks[i];
        seg_info.seg_type = region_classes[i];
        one_road.vec_road.push_back(seg_info);
    }
    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;
    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 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) {
        SegInfo seg_info;
        seg_info.seg_type = region_classes[i];
        seg_info.vec_pt = poly_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);
    }
    // cv::imwrite("./direct.jpg", image);
    // image.release();
    return vec_road;
}
int RoadSegAnalysis::release() {
    ACL_CALL(aclrtSetCurrentContext(ctx), SY_SUCCESS, SY_FAILED);	
    if (m_handle) {
        rs_release(&m_handle);
    }
    if(ctx){
        aclrtDestroyContext(ctx);
        ctx = nullptr;
    }
    return SY_SUCCESS;
}
// head_or_tail 0：车头  1：车尾
int RoadSegAnalysis::check_reverse_driving(std::vector<SegInfo>& vec_direct, sy_rect src_rc, int src_width, int src_height, int head_or_tail) {
    float scale_w = 640.0 / src_width;
	float scale_h = 360.0 / src_height;
    sy_rect rc;
    rc.left_ = src_rc.left_ * scale_w;
    rc.width_ = src_rc.width_ * scale_w;
    rc.top_ = src_rc.top_ * scale_h;
    rc.height_ = src_rc.height_ * scale_h;
    std::vector<cv::Point> polygon_pts;
    cv::Point pt_lt;
    pt_lt.x = rc.left_;
    pt_lt.y = rc.top_;
    polygon_pts.push_back(pt_lt);
    cv::Point pt_rt;
    pt_rt.x = rc.left_ + rc.width_;
    pt_rt.y = rc.top_;
    polygon_pts.push_back(pt_rt);
    cv::Point pt_rb;
    pt_rb.x = rc.left_ + rc.width_;
    pt_rb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_rb);
    cv::Point pt_lb;
    pt_lb.x = rc.left_;
    pt_lb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_lb);
    int coming_count = 0;
    int leaving_count = 0;
    for (size_t i = 0; i < vec_direct.size(); i++) {
        SegInfo& region = vec_direct[i];
        if (region.seg_type == 1) {
            for (size_t j = 0; j < polygon_pts.size(); j++) {
                double dist  = pointPolygonTest(region.vec_pt, polygon_pts[j], false);
                if (dist > 0) {
                    coming_count ++;
                }
            }
        } else if (region.seg_type == 2) {
            for (size_t j = 0; j < polygon_pts.size(); j++) {
                double dist  = pointPolygonTest(region.vec_pt, polygon_pts[j], false);
                if (dist > 0) {
                    leaving_count ++;
                }
            }
        }
    }
    int direct = 0;
    if(coming_count > 0 && leaving_count > 0) {
        if (coming_count > leaving_count) {
            direct = 1;
        } else {
            direct = 2;
        }
    } else if(coming_count > 0 && leaving_count <= 0) {
        direct = 1;
    } else if(coming_count <= 0 && leaving_count > 0) {
        direct = 2;
    } 
    // int center_x = (rc.left_ + rc.width_ / 2) * scale_w;
    // int center_y = (rc.top_ + rc.height_ / 2) * scale_h;
    // if(center_x < 0 ||center_x >= direct_mask.cols || center_y < 0 || center_y >= direct_mask.rows){
    //     return -1;
    // }
    // '来': 1,  '去': 2,  '近': 3,  '远': 4
    // int direct = direct_mask.at<int>(center_x, center_y);
    if (direct == 1 && head_or_tail == 0) {
        // 来车道，车头，正常行驶
        return 0;
    } else if (direct == 1 && head_or_tail == 1){
        // 来车道，车尾，逆行
        return 1;
    } else if (direct == 2 && head_or_tail == 0){
        // 去车道，车头，逆行
        return 1;
    } else if (direct == 2 && head_or_tail == 1){
        // 去车道，车尾，正常行驶
        return 0;
    } 
    
    return -1;
}
int RoadSegAnalysis::check_cross_line(std::vector<LineInfo>& vec_line, sy_rect src_rc, int src_width, int src_height) {
    float scale_w = 640.0 / src_width;
    float scale_h = 360.0 / src_height;
    sy_rect rc;
    rc.left_ = src_rc.left_ * scale_w;
    rc.width_ = src_rc.width_ * scale_w;
    rc.top_ = src_rc.top_ * scale_h;
    rc.height_ = src_rc.height_ * scale_h;
    std::vector<cv::Point> polygon_pts;
    cv::Point pt_lt;
    pt_lt.x = rc.left_;
    pt_lt.y = rc.top_;
    polygon_pts.push_back(pt_lt);
    cv::Point pt_rt;
    pt_rt.x = rc.left_ + rc.width_;
    pt_rt.y = rc.top_;
    polygon_pts.push_back(pt_rt);
    cv::Point pt_rb;
    pt_rb.x = rc.left_ + rc.width_;
    pt_rb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_rb);
    cv::Point pt_lb;
    pt_lb.x = rc.left_;
    pt_lb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_lb);
    for (size_t i = 0; i < vec_line.size(); i++) {
        LineInfo& line = vec_line[i];
        if (line.line_type == 1 || line.line_type == 2) {
            // 黄实线
            int in_count = 0;
            for (size_t j = 0; j < line.vec_pt.size(); j++) {
                double dist  = pointPolygonTest(polygon_pts, line.vec_pt[j], false);
                if (dist > 0) {
                    in_count ++;
                }
            }
            if (in_count > 5)
            {//有5个点就认为是压线了
                return line.line_type;
            }
        }
    }
    
    return -1;
}
int RoadSegAnalysis::check_cross_region(std::vector<SegInfo>& vec_reg, sy_rect src_rc, int src_width, int src_height, int region_type) {
    float scale_w = 640.0 / src_width;
    float scale_h = 360.0 / src_height;
    sy_rect rc;
    rc.left_ = src_rc.left_ * scale_w;
    rc.width_ = src_rc.width_ * scale_w;
    rc.top_ = src_rc.top_ * scale_h;
    rc.height_ = src_rc.height_ * scale_h;
    std::vector<cv::Point> polygon_pts;
    cv::Point pt_lt;
    pt_lt.x = rc.left_;
    pt_lt.y = rc.top_;
    polygon_pts.push_back(pt_lt);
    cv::Point pt_rt;
    pt_rt.x = rc.left_ + rc.width_;
    pt_rt.y = rc.top_;
    polygon_pts.push_back(pt_rt);
    cv::Point pt_rb;
    pt_rb.x = rc.left_ + rc.width_;
    pt_rb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_rb);
    cv::Point pt_lb;
    pt_lb.x = rc.left_;
    pt_lb.y = rc.top_ + rc.height_;
    polygon_pts.push_back(pt_lb);
    for (size_t i = 0; i < vec_reg.size(); i++) {
        SegInfo& seg = vec_reg[i];
        if (seg.seg_type == region_type) {
            
            int in_count = 0;
            // 车辆与region相交
            for (size_t j = 0; j < seg.vec_pt.size(); j++) {
                double dist  = pointPolygonTest(polygon_pts, seg.vec_pt[j], false);
                if (dist > 0) {
                    in_count ++;
                }
            }
            if (in_count > 5)
            {//有5个点就认为是压线了
                return 1;
            }
            // 车辆与region相交情形未检测出来，检测车辆在region中情形
            for (size_t j = 0; j < polygon_pts.size(); j++) {
                double dist  = pointPolygonTest(seg.vec_pt, polygon_pts[j], false);
                if (dist > 0) {
                    return 1;
                }
            }
        }
    }
    
    return -1;
}
cv::Mat RoadSegAnalysis::mask_to_rgb(cv::Mat img, cv::Mat mask) {
    cv::Mat masks = img.clone();
    int reg_cls = 8;
    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;
}
float RoadSegAnalysis::contourArea(std::vector<cv::Point> contour, cv::Point2f& center) {
    cv::RotatedRect rect = cv::minAreaRect(contour);  // 最小外接矩形 rect[0]中心点 rect[1]宽 高  rect[2]旋转角度
    center = rect.center;
    return cv::contourArea(contour);
}
void RoadSegAnalysis::lanes_process(const rs_lane* lanes, int lane_count, std::vector<std::pair<std::vector<cv::Point>, int>>& combined, float scale_w, float scale_h) {   
    std::vector<std::vector<cv::Point> > lanes_xys;
    std::vector<int> lanes_cls;
	for (int i = 0; i < lane_count; i++) {
        std::vector<cv::Point> xys;
		for (int j = 0; j < lanes[i].num_points; j++) {
            int x = static_cast<int>(lanes[i].points[j].x_ * scale_w);
            int y = static_cast<int>(lanes[i].points[j].y_ * scale_h);
            if (x > 0 && y > 0) {
                xys.emplace_back(x, y);
            }
        }
        if (!xys.empty()) {
            lanes_xys.push_back(xys);
            lanes_cls.push_back(lanes[i].cls);
        }
    }
    for (size_t i = 0; i < lanes_xys.size(); ++i) {
        combined.push_back(std::make_pair(lanes_xys[i], lanes_cls[i]));
    }
    if (!combined.empty()) {
        //按车道线起点坐标排序,相应的类别顺序也会变化以保证标签对齐
		std::sort(combined.begin(), combined.end(), [](const std::pair<std::vector<cv::Point>, int>& a, const std::pair<std::vector<cv::Point>, int>& b) {
			return a.first[0].x < b.first[0].x;
		});
    }
}
cv::Mat RoadSegAnalysis::imshow_lanes(cv::Mat img, const rs_lane* lanes, int lane_count) {
	float scale_w = img.cols / 640.0;
	float scale_h = img.rows / 360.0;
    std::vector<std::pair<std::vector<cv::Point>, int>> combined;
    lanes_process(lanes, lane_count, combined, scale_w, scale_h);
    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
		cv::Scalar color(lane_colors[cls][0],lane_colors[cls][1],lane_colors[cls][2]);
        for (size_t i = 1; i < xys.size(); ++i) {
            cv::line(img, xys[i - 1], xys[i], color, 4);
        }
    }
    
    return img;
}
int RoadSegAnalysis::Mask2LanePoints(const cv::Mat& pred, std::vector<std::vector<cv::Point>>&lanes, std::vector<int>& cats) {
    std::vector<int> labels = {8, 9, 10, 11};
    for(auto cat: labels) {
        cv::Mat b_masks, measure_labels, stats, centroids;
        cv::compare(pred, cat, b_masks, cv::CMP_EQ); //将pred元素逐个和cat比较，相同255，不同0
        // cv::threshold(pred, b_masks, 126, 255, cv::THRESH_OTSU);    //生成二值图
        // 连通域计算 b_masks:闭操作后的二值图像 measure_labels:和原图一样大的标记图 centroids:nccomps×2的矩阵,表示每个连通域的质心(x, y)
        // stats:nccomps×5的矩阵 表示每个连通区域的外接矩形和面积(x, y, w, h, area)，索引0是背景信息
        int nccomps = cv::connectedComponentsWithStats (b_masks, measure_labels, stats, centroids, 4); //8
        for(int cv_measure_id = 1; cv_measure_id < nccomps; cv_measure_id++ ) {	//跳过背景信息0
            cv::Mat cv_measure_mask;
            cv::compare(measure_labels, cv_measure_id, cv_measure_mask, cv::CMP_EQ);
            std::vector< std::vector< cv::Point> > contours;
            cv::findContours(cv_measure_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++) {
                float area = cv::contourArea(contours[j]);  if (area < 60) continue; //30
                cv::approxPolyDP(cv::Mat(contours[j]), contours_poly, 6, true); //减小epsilon提高拟合精度(需要调试，epsilon过小会使多边形不够简化，单条线变多条)
                if (contours_poly.size() == 1) continue;
                lanes.push_back(contours_poly);
                cats.push_back(cat);
            }
        }    
    }
    return 0;
}
cv::Mat RoadSegAnalysis::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) {
    std::vector<int> pred_cls; 
    int h = 360, w = 640;
    cv::Mat lanes_masks = cv::Mat_<int>(h,w); //正常分割结果
    cv::Mat mask_rmlane = cv::Mat_<int>(h,w); //车道线区域置为背景
    cv::Mat solve_masks = cv::Mat_<int>(h,w); //计算主行驶区域用(虚线及减速标线归入行车道)
    int step_size = h*w;
    int seg_min_region_area = 512; //1024 
	for (int i = 0; i < h; ++i) {
		for (int j = 0; j < w; ++j) {
			int max_cls = seg_array[(i * w + j)];
			lanes_masks.ptr<int>(i)[j] = max_cls;
            pred_cls.push_back(max_cls);
            mask_rmlane.ptr<int>(i)[j] = max_cls;
            solve_masks.ptr<int>(i)[j] = max_cls;
		}
	}
    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
		if (cls == 1) cls = 8;  /*黄实线*/ if (cls == 2) cls = 9;  /*白实线*/
        if (cls == 3) cls = 10; /*虚线*/   if (cls == 4) cls = 11; /*黄虚线*/
        for (size_t i = 1; i < xys.size(); ++i) {
            cv::line(lanes_masks, xys[i - 1], xys[i], cls, 4); //绘制车道线用于求连通域
            if (cls == 10 || cls == 11)   cls = 1; //求主行驶区域时将虚线归入行车道 
			cv::line(solve_masks, xys[i - 1], xys[i], cls, 4);
            
        }
    }
    /* 求背景区域--mask车道区域，场景变化用
    cv::compare(lanes_masks, 0, background_mask, cv::CMP_EQ); //将lanes_masks元素逐个和0比较，相同255，不同0*/
#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
    //2.去重获取预测到的类别
    std::vector<int> labels(pred_cls);
    std::sort(labels.begin(),labels.end());
    labels.erase(std::unique(labels.begin(),labels.end()),labels.end());
  
    //4.求车道线区域
    int flag = Mask2LanePoints(lanes_masks, lanes, cats); ///////////////////////// 车道线如何与mask结合
 
    //5.求道路区域
    int count = 0;
    for(auto cat: labels) {
        // std::cout << cat << std::endl;
        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
            //移除过小的区域，并将对应位置置为0
            if (stats.at<int>(i, cv::CC_STAT_AREA) < seg_min_region_area) {
                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;
            }
            double centr_x = centroids.at<double>(i, 0);
            double centr_y = centroids.at<double>(i, 1);
            int region_class = cat;
            // printf("region_class: %d\n", region_class);
            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);
                if (x_sort.count(centr_x)) centr_x += 0.0001;
                x_sort.insert(std::make_pair(centr_x, count));
                ++ count;
            }
        }    
    }
#if 0
    //6.draw lanes
    for (int i = 0; i < lanes.size(); ++i) {
        int thickness = 4;
        for (int j = 0; j < lanes[i].size()-1; ++j) {
            cv::line(mask_rmlane, lanes[i][j], lanes[i][j+1], {seg_colors[cats[i]][0],seg_colors[cats[i]][1],seg_colors[cats[i]][2]}, thickness);
            // std::cout << lanes[i][j] << " " << lanes[i][j+1] << " " << cats[i] << std::endl;
        }
    }
#endif
    return mask_rmlane;
}