diff --git a/PINetTensorrt.cpp b/PINetTensorrt.cpp index 1af81aa..e6f6f43 100644 --- a/PINetTensorrt.cpp +++ b/PINetTensorrt.cpp @@ -389,7 +389,7 @@ void PINetTensorrt::generatePostData(float* confidance_data, float* offsets_data for (int j = 0; j < dim.d[2]; ++j) { if ((int)mask.at(i, j)) { cv::Vec2f pointOffset = offsets.at(i, j); - cv::Point2f point(pointOffset[1] + j, pointOffset[0] + i); + cv::Point2f point(pointOffset[0] + j, pointOffset[1] + i); cv::circle(offsetImage, point * 8, 3, color, -1); } } @@ -417,14 +417,19 @@ LaneLines PINetTensorrt::generateLaneLine(float* confidance_data, float* offsets LaneLines laneLines; std::vector laneFeatures; - auto findNearestFeature = [&laneFeatures](const cv::Vec4f& feature) -> int { + + auto findNearestFeature = [&laneFeatures](const cv::Vec4f& feature) -> std::pair { + int index=-1; + float min_feature_dis = 10000.; for (int i = 0; i < laneFeatures.size(); ++i) { auto delta = laneFeatures[i] - feature; - if (delta.dot(delta) <= threshold_instance) { - return i; + auto alpha = pow(delta[0],2)+pow(delta[1],2)+pow(delta[2],2)+pow(delta[3],2); + if (sqrt(alpha) <= min_feature_dis){ + index = i; + min_feature_dis = sqrt(alpha); } } - return -1; + return std::pair(index, min_feature_dis); }; for (int i = 0; i < dim.d[1]; ++i) { @@ -434,25 +439,31 @@ LaneLines PINetTensorrt::generateLaneLine(float* confidance_data, float* offsets } const cv::Vec2f& offset = offsets.at(i, j); - cv::Point2f point(offset[1] + j, offset[0] + i); + cv::Point2f point(offset[0] + j, offset[1] + i); if (point.x > dim.d[2] || point.x < 0.f) continue; if (point.y > dim.d[1] || point.y < 0.f) continue; const cv::Vec4f& feature = features.at(i, j); - int lane_index = findNearestFeature(feature); - - if (lane_index == -1) { + std::pair lane_index = findNearestFeature(feature); + + if (lane_index.first == -1) { + laneLines.emplace_back(LaneLine({point})); + laneFeatures.emplace_back(feature); + } + else if (lane_index.second <= threshold_instance ) { + + auto& laneline = laneLines[lane_index.first ]; + auto& lanefeature = laneFeatures[lane_index.first ]; + if (lane_index.second <= threshold_instance ){ + auto point_size = laneline.size(); + lanefeature = lanefeature.mul(cv::Vec4f::all(point_size)) + feature; + lanefeature = lanefeature.mul(cv::Vec4f::all(1.f / (point_size + 1))); + laneline.emplace_back(point); + } + } + else{ laneLines.emplace_back(LaneLine({point})); laneFeatures.emplace_back(feature); - } else { - auto& laneline = laneLines[lane_index]; - auto& lanefeature = laneFeatures[lane_index]; - - auto point_size = laneline.size(); - - lanefeature = lanefeature.mul(cv::Vec4f::all(point_size)) + feature; - lanefeature = lanefeature.mul(cv::Vec4f::all(1.f / (point_size + 1))); - laneline.emplace_back(point); } } }