#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}};

void lanes_process(const rs_lane* lanes, int lane_count, std::vector<std::pair<std::vector<cv::Point>, int>>& combined, float scale_w = 1.0, float scale_h = 1.0) {   
    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 mask_to_rgb(cv::Mat img, cv::Mat mask) {
    cv::Mat rgb = img.clone();
    int reg_cls = 9;
    for (int i = 0; i < rgb.rows; i++) {
        for (int j = 0; j < rgb.cols; j++) {
            for (int k = 1; k < reg_cls; k++) {
                if (mask.at<int>(i,j) == k) {
                    rgb.at<cv::Vec3b>(i,j)[0] = seg_colors[k][0];
                    rgb.at<cv::Vec3b>(i,j)[1] = seg_colors[k][1];
                    rgb.at<cv::Vec3b>(i,j)[2] = seg_colors[k][2];
                }
            }
        }
    }
    return rgb;
}

float 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);
}

int Mask2LanePoints(const cv::Mat& pred, std::vector<std::vector<cv::Point>>&lanes, std::vector<int>& cats) {
    std::vector<int> labels = {9, 10, 11, 12, 13};
    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
                float area = cv::contourArea(contours[j]);  if (area < 30) 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;
}

int findMainMaskV2(bool large_resolution, const cv::Mat& solve_masks, const std::vector<int>& labels, std::vector<std::vector<cv::Point>> &poly_masks, std::vector<int> &region_classes, std::vector<std::vector<cv::Point>> &main_masks, cv::Mat& new_mask, cv::Mat& new_mask_forlane) {
    int h = 360, w = 640;
	int seg_num_cls = 8;
    int seg_num_seg = 4;
	int seg_min_region_area = 512; //10
     //  mask部分区域
    if (large_resolution)   solve_masks(cv::Rect(0, 0, w, int(h * 0.14))) = 0;   
    else    solve_masks(cv::Rect(0, 0, w, int(h * 0.22))) = 0;           
    solve_masks(cv::Rect(0, int(h * 0.95), w, int(h * 0.05))) = 0;
    solve_masks(cv::Rect(0, 0, int(w * 0.02), h)) = 0;
    solve_masks(cv::Rect(int(w * 0.95), 0, int(w * 0.05), h)) = 0;
   
    float hf_w = w*0.5;
    std::vector<std::vector<cv::Point>> all_emerg_contours, all_com_contours, mid_com_contours, mid_zd_contours;
    std::vector<cv::Point> emerg_contours, mid_contours, com_contours, zd_contours;
    float emerg_max_area = 0, max_area = 0, zd_max_area = 0;
    
    for(auto cat: labels) {
        if (cat == 0 || cat > seg_num_seg) 
            continue;

        // 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(solve_masks, cat, b_masks, cv::CMP_EQ); //将solve_masks元素逐个和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(solve_masks,n_masks,solve_masks);
                continue;
            }

            double centr_x = centroids.at<double>(i, 0);
            double centr_y = centroids.at<double>(i, 1);

            int region_class = cat;
            cv::Mat region_mask;
            cv::bitwise_and(measure_labels,n_masks,measure_labels);/////
            cv::compare(measure_labels, i, region_mask, cv::CMP_EQ);

            // 闭运算，先膨胀再腐蚀，去除较暗部分
            cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(17, 5));
            cv::morphologyEx(region_mask, region_mask, cv::MORPH_CLOSE, kernel);

            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;

            // 应急车道
            if(region_class == 2) {
                for (auto contour : contours) {
                    cv::Point2f center; 
                    float area_ = contourArea(contour, center);
                    if (area_ > emerg_max_area) {
                        emerg_max_area = area_;
                        emerg_contours = contour;
                    }  
                    all_emerg_contours.push_back(contour);
                }       
            }

            // 普通行车道
            if(region_class == 1) {
                for (auto contour : contours) {
                    cv::Point2f center; 
                    float area_ = contourArea(contour, center);
                    if (area_ > max_area) {
                        max_area = area_;
                        com_contours = contour;
                    } 
                    // 取中点所在的连通域
                    cv::Point2f point(int(w*0.5),int(h*0.95)-1);
                    double distance = cv::pointPolygonTest(contour, point, false);
                    if (distance >= 0) {
                        mid_contours = contour;
                        hf_w = center.x;
                        continue;
                    }
                    if (abs(center.x - hf_w) < 90) // 取中点所在区域左右90像素内的连通域 230307
                        mid_com_contours.push_back(contour);
                }       
            }
            
            // 导流线
            if(region_class == 3) {
                for (auto contour : contours) {
                    cv::Point2f center; 
                    float area_ = contourArea(contour, center);
                    if (area_ > zd_max_area) {
                        zd_max_area = area_;
                        zd_contours = contour;
                    } 
                    // if (abs(center.x - w*0.5) < 200) // 取中点左右200像素内的连通域 
					if (abs(center.x - w*0.5) < 250) // 取中点左右250像素内的连通域 
                        mid_zd_contours.push_back(contour);
                }       
            }
         
            // 拟合各行驶区域
            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);
            }
        }
    }

    // cv::Mat new_mask(h, w, CV_8UC3, cv::Scalar(0,0,0));
    if (!mid_contours.empty()) {
        // cv::fillConvexPoly(new_mask, mid_contours, cv::Scalar(255,255,255)); //填充255
        cv::fillConvexPoly(new_mask, mid_contours, cv::Scalar(255)); //填充255
        cv::fillConvexPoly(new_mask_forlane, mid_contours, cv::Scalar(255)); //填充255
        std::vector<cv::Point> contours_poly;
        cv::approxPolyDP(cv::Mat(mid_contours), contours_poly, 10, true);
        if (contours_poly.size() > 2) main_masks.push_back(contours_poly);
        all_com_contours.push_back(mid_contours);
    } 
    else {
        // cv::fillConvexPoly(new_mask, com_contours, cv::Scalar(255,255,255)); //填充255
        cv::fillConvexPoly(new_mask, com_contours, cv::Scalar(255)); //填充255
        cv::fillConvexPoly(new_mask_forlane, com_contours, cv::Scalar(255)); //填充255
        std::vector<cv::Point> contours_poly;
        cv::approxPolyDP(cv::Mat(com_contours), contours_poly, 10, true);
        if (contours_poly.size() > 2) main_masks.push_back(contours_poly);
        all_com_contours.push_back(com_contours);
    } 

    for (auto contour_ : mid_com_contours) {
        // cv::fillConvexPoly(new_mask, contour_, cv::Scalar(255,255,255)); //填充255
        cv::fillConvexPoly(new_mask, contour_, cv::Scalar(255)); //填充255
        cv::fillConvexPoly(new_mask_forlane, contour_, cv::Scalar(255)); //填充255
        std::vector<cv::Point> contours_poly;
        cv::approxPolyDP(cv::Mat(contour_), contours_poly, 10, true);
        if (contours_poly.size() > 2) main_masks.push_back(contours_poly);
        all_com_contours.push_back(contour_);
    }
   
    if (!emerg_contours.empty() && !all_emerg_contours.empty()) {
        if (all_emerg_contours.size() == 1) 
            // cv::fillConvexPoly(new_mask, emerg_contours, cv::Scalar(255,255,255));
            cv::fillConvexPoly(new_mask, emerg_contours, cv::Scalar(255)); //填充255
        else {
            float sum_centerx = 0, sum_centery = 0;
            float length_emerg_com = 9999999;
            // 计算选中行车道的平均中心点
            for (auto contour_ : all_com_contours) {
                cv::Point2f center; 
                float area_ = contourArea(contour_, center);
                sum_centerx += center.x;
                sum_centery += center.y;
            }
            float ave_centerx = sum_centerx / all_com_contours.size();
            float ave_centery = sum_centery / all_com_contours.size();
            // 选择距离最近的应急车道且面积不小于最大应急车道的一半
            for (auto econtour_ : all_emerg_contours) {
                cv::Point2f center; 
                float area_ = contourArea(econtour_, center); 
                float tmp_length = sqrt(pow((ave_centerx-center.x),2.0) + pow((ave_centery-center.y),2.0));
                if (tmp_length < length_emerg_com and area_ > 0.7*emerg_max_area) {
                    length_emerg_com = tmp_length;
                    emerg_contours = econtour_;
                }
            }
            // cv::fillConvexPoly(new_mask, emerg_contours, cv::Scalar(255,255,255));  
            cv::fillConvexPoly(new_mask, emerg_contours, cv::Scalar(255)); //填充255     
        }
    }

    if (!mid_zd_contours.empty()) {
        for (auto contour_ : mid_zd_contours) 
            // cv::fillConvexPoly(new_mask, contour_, cv::Scalar(255,255,255)); //填充255 
            cv::fillConvexPoly(new_mask, contour_, cv::Scalar(255)); //填充255
    }

    // 膨胀腐蚀，去除行车道和应急车道间的空洞
    cv::Mat kernel_ = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(17, 5));
    cv::morphologyEx(new_mask, new_mask, cv::MORPH_CLOSE, kernel_);
    cv::morphologyEx(new_mask_forlane, new_mask_forlane, cv::MORPH_CLOSE, kernel_);

    return 0;
}


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 &background_mask) {
    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 = 9;  /*黄实线*/ if (cls == 2) cls = 10;  /*白实线*/
        if (cls == 3) cls = 11; /*虚线*/   if (cls == 4) cls = 12; /*黄虚线*/
        if (cls == 5) cls = 13; /*车道中线*/
        for (size_t i = 1; i < xys.size(); ++i) {
            cv::line(lanes_masks, xys[i - 1], xys[i], cls, 4); //绘制车道线用于求连通域
            if (cls == 11 || cls == 12)   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 0
    //--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;
                // for (auto iter: contours_poly) {
                //     std::cout << "contour " << iter << " " << cat << std::endl;
                // }
    
            }
        }    
    }

#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], cats[i], thickness);
            // std::cout << lanes[i][j] << " " << lanes[i][j+1] << " " << cats[i] << std::endl;
        }
    }
#endif
    return mask_rmlane;
}



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) {
    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;
}


// 合并来去区域
cv::Mat merge_direct_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)];
            if (max_cls == 1 || max_cls == 2) max_cls = 5; // 来去区域合并的类别为5
            if (max_cls == 3 || max_cls == 4) max_cls = 6; // 近远区域合并的类别为6
            pred_cls.push_back(max_cls);
            mask_rmlane.ptr<int>(i)[j] = max_cls;
		}
	}

#if 0 /*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;
}

RoadSegAnalysis::RoadSegAnalysis(/* args */)
{
}

RoadSegAnalysis::~RoadSegAnalysis()
{
    release();
}

int RoadSegAnalysis::init(int devId){
    ACL_CALL(aclrtCreateContext(&ctx, devId), SY_SUCCESS, SY_FAILED);

    rs_param param;
    param.modelNames = "./models/road_seg/tzroad_seg240108_310p.om";
    param.thresld = 0.25;
    param.devId = devId;

    cout << "rs_init start " << endl;
    int ret = rs_init(&m_handle, param);
    if (ret != 0) {
        return -1;
	}

    cout << "rs_init success " << endl;

    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;
            for (size_t i = 0; i < one_result.lane_count; i++)
            {
                auto one_lane = one_result.reg_array[i];
                LineInfo info;
                for (size_t j = 0; j < one_lane.num_points; j++)
                {
                    info.vec_pt.push_back(one_lane.points[j]);
                }
                info.line_type = one_lane.cls;
                one_road.vec_line.push_back(info);
            }

            one_road.vec_road = parse_seg(one_result, vec_img[b]);
            one_road.direct_mask = parse_direct(one_result, vec_img[b]);

            post_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;
}

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;
}

std::vector<SegInfo> RoadSegAnalysis::parse_seg(rs_result one_result, sy_img src) {
    int w = 640;
    int h = 360;
    float alpha = 0.75;
    cv::Mat overlayed_img(cv::Size(w,h), CV_8UC3, cv::Scalar(0, 0, 0));
    // 将车道线标签转换为彩色图像
    for (int i = 0; i < h; ++i) {
        for (int j = 0; j < w; ++j) {
            int idx = one_result.seg_array[(i * w + j)];
            overlayed_img.at<cv::Vec3b>(i, j)[0] = seg_colors[idx][0]; // R通道;
            overlayed_img.at<cv::Vec3b>(i, j)[1] = seg_colors[idx][1]; // G通道
            overlayed_img.at<cv::Vec3b>(i, j)[2] = seg_colors[idx][2]; // B通道
        }
    }

    cv::resize(overlayed_img, overlayed_img, cv::Size(src.w_,src.h_), 0, 0, cv::INTER_LINEAR); 
    // 将原始图像和彩色车道线图进行混合
    // cv::addWeighted(cvImg, alpha, overlayed_img, 1 - alpha, 0, overlayed_img);

    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<SegInfo> vec_seg;
    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
        SegInfo info;
        info.seg_type = cls;
		// cv::Scalar color(lane_colors[cls][0],lane_colors[cls][1],lane_colors[cls][2]);
        for (size_t i = 0; i < xys.size(); ++i) {
            // cv::line(overlayed_img, xys[i - 1], xys[i], color, 4);
            sy_point pt;
            pt.x_ = xys[i].x;
            pt.y_ = xys[i].y;
            info.vec_pt.push_back(pt);
        }
        vec_seg.push_back(info);
    }

    return vec_seg;
}

cv::Mat RoadSegAnalysis::parse_direct(rs_result one_result, sy_img src){
    int src_width = src.w_;

    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, direct_masks, merge_masks;
    std::vector<int> region_classes, cats, direct_classes, merge_classes;
    cv::Mat background_mask;
    std::map<double, int> x_sort;
    bool large_resolution = false;
    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, background_mask); //m_masks:mask前的结果  poly_masks后的结果
    cv::Mat direct_output = direct_post_process(large_resolution, one_result.direct_seg, direct_masks, direct_classes);


    // std::vector<std::vector<cv::Point>> direct_masks;
    // std::vector<int> direct_classes;
    // bool large_resolution = false;
    // if (src_width > 1920) large_resolution = true;
    // 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 merge_image = mask_to_rgb(image, direct_output);
    // cv::imwrite("direct.jpg", merge_image);

    int direct = direct_output.at<int>(455, 146);

    return direct_output;
}

// head_or_tail 0：车头  1：车尾
int RoadSegAnalysis::check_reverse_driving(cv::Mat direct_mask, sy_rect rc, int src_width, int src_height, int head_or_tail) {
    // cv::Mat image(cv::Size(640, 360), CV_8UC3, cv::Scalar(0, 0, 0));
    // cv::Mat merge_image = mask_to_rgb(image, direct_mask);
    // cv::imwrite("direct_mask.jpg", merge_image);

    float scale_w = 640.0 / src_width;
	float scale_h = 360.0 / src_height;
    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 rc) {

    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;
            std::vector<cv::Point> vec_pt;
            for (size_t j = 0; j < line.vec_pt.size(); j++) {
                double dist  = pointPolygonTest(polygon_pts, cv::Point2f(line.vec_pt[j].x_, line.vec_pt[j].y_), 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 rc, int region_type) {

    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;
            std::vector<cv::Point> vec_pt;
            for (size_t j = 0; j < seg.vec_pt.size(); j++) {
                double dist  = pointPolygonTest(polygon_pts, cv::Point2f(seg.vec_pt[j].x_, seg.vec_pt[j].y_), false);
                if (dist > 0) {
                    in_count ++;
                }
            }

            if (in_count > 5)
            {//有5个点就认为是压线了
                return 1;
            }
        }
    }
    
    return -1;
}

cv::Mat 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;
}

void RoadSegAnalysis::post_direct(rs_result one_result, sy_img src){
	int src_width = src.w_;
	int src_height = src.h_;
    int w = 640;
    int h = 360;
    float alpha = 0.75;
    cv::Mat overlayed_direct_img(cv::Size(w,h), CV_8UC3, cv::Scalar(0, 0, 0));
    // 将车道线标签转换为彩色图像
    for (int i = 0; i < h; ++i) {
        for (int j = 0; j < w; ++j) {
            int idx = one_result.direct_seg[(i * w + j)];
            overlayed_direct_img.at<cv::Vec3b>(i, j)[0] = seg_colors[idx][0]; // R通道;
            overlayed_direct_img.at<cv::Vec3b>(i, j)[1] = seg_colors[idx][1]; // G通道
            overlayed_direct_img.at<cv::Vec3b>(i, j)[2] = seg_colors[idx][2]; // B通道
        }
    }
    cv::resize(overlayed_direct_img, overlayed_direct_img, cv::Size(src_width,src_height), 0, 0, cv::INTER_LINEAR); 
    // 将原始图像和彩色车道线图进行混合
    // cv::addWeighted(cvImg, alpha, overlayed_direct_img, 1 - alpha, 0, overlayed_direct_img);
    cv::Mat img_direct_lane = imshow_lanes(overlayed_direct_img, one_result.reg_array, one_result.lane_count);
    cv::imwrite("img_direct_lane.jpg", img_direct_lane);


    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, direct_masks, merge_masks;
    std::vector<int> region_classes, cats, direct_classes, merge_classes;
    cv::Mat background_mask;
    std::map<double, int> x_sort;
    bool large_resolution = false;
    if (src_height > 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, background_mask); //m_masks:mask前的结果  poly_masks后的结果
    cv::Mat direct_output = direct_post_process(large_resolution, one_result.direct_seg, direct_masks, direct_classes);
    cv::Mat merge_output = merge_direct_process(large_resolution, one_result.direct_seg, merge_masks, merge_classes);

    cv::Mat image(cv::Size(w,h), CV_8UC3, cv::Scalar(0, 0, 0));
    cv::Mat vis_image = mask_to_rgb(image, seg_output);
    cv::Mat direct_image = mask_to_rgb(image, direct_output);
    cv::Mat merge_image = mask_to_rgb(image, merge_output);
    cv::imwrite("vis_image.jpg", vis_image);
    cv::imwrite("direct_image.jpg", direct_image);
    cv::imwrite("merge_image.jpg", merge_image);
    cv::imwrite("background_mask.jpg", background_mask);
}