Line data Source code
1 : #include "../../includes/inference/LanePostProcessor.hpp"
2 : #include <opencv2/cudaarithm.hpp>
3 : #include <cmath>
4 : #include <numeric>
5 :
6 23 : LanePostProcessor::LanePostProcessor(int minArea, int minLength, float angleThresh, float mergeDist)
7 : : minComponentSize(minArea),
8 : minComponentLength(minLength),
9 : angleThreshold(angleThresh),
10 23 : mergeDistance(mergeDist) {}
11 :
12 0 : cv::cuda::GpuMat LanePostProcessor::process(const cv::cuda::GpuMat& rawMaskGpu) {
13 0 : cv::cuda::GpuMat binaryMaskGpu;
14 0 : cv::cuda::threshold(rawMaskGpu, binaryMaskGpu, 0.5, 255.0, cv::THRESH_BINARY);
15 0 : binaryMaskGpu.convertTo(binaryMaskGpu, CV_8U);
16 :
17 0 : cv::Mat binaryMask;
18 0 : binaryMaskGpu.download(binaryMask); // CPU copy for contour analysis
19 :
20 0 : std::vector<LaneInfo> initialLanes = extractLaneInfo(binaryMask, false, true);
21 :
22 0 : cv::Mat mergedMask;
23 0 : cv::cvtColor(binaryMask, mergedMask, cv::COLOR_GRAY2BGR);
24 0 : drawLaneConnections(initialLanes, mergedMask);
25 :
26 0 : cv::Mat grayMerged;
27 0 : cv::cvtColor(mergedMask, grayMerged, cv::COLOR_BGR2GRAY);
28 :
29 0 : std::vector<LaneInfo> finalLanes = extractLaneInfo(grayMerged, true, false);
30 0 : cv::Mat filtered = renderFilteredMask(finalLanes, binaryMask.size());
31 :
32 0 : cv::cuda::GpuMat resultGpu;
33 0 : resultGpu.upload(filtered);
34 0 : return resultGpu;
35 : }
36 :
37 0 : std::vector<LanePostProcessor::LaneInfo> LanePostProcessor::extractLaneInfo(const cv::Mat& mask, bool filterSize, bool filterLength) {
38 0 : std::vector<std::vector<cv::Point>> contours;
39 0 : cv::findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
40 :
41 0 : std::vector<LaneInfo> lanes;
42 0 : int height = mask.rows, width = mask.cols;
43 :
44 0 : for (const auto& cnt : contours) {
45 0 : double area = cv::contourArea(cnt);
46 0 : double length = cv::arcLength(cnt, false);
47 :
48 0 : if ((!filterSize || area > minComponentSize) &&
49 0 : (!filterLength || length > minComponentLength) &&
50 0 : cnt.size() >= 2) {
51 :
52 0 : cv::Vec4f line;
53 0 : cv::fitLine(cnt, line, cv::DIST_L2, 0, 0.01, 0.01);
54 0 : float vx = line[0], vy = line[1], x0 = line[2], y0 = line[3];
55 0 : float angle = std::fmod(std::atan2(vy, vx) * 180.0 / CV_PI, 180.0f);
56 :
57 0 : cv::Moments M = cv::moments(cnt);
58 0 : if (M.m00 == 0) continue;
59 0 : cv::Point centroid(M.m10 / M.m00, M.m01 / M.m00);
60 :
61 0 : std::vector<cv::Point> intersections;
62 0 : if (vy != 0) {
63 0 : float t = -y0 / vy;
64 0 : float x = x0 + vx * t;
65 0 : if (x >= 0 && x < width) intersections.emplace_back(x, 0);
66 0 : t = (height - 1 - y0) / vy;
67 0 : x = x0 + vx * t;
68 0 : if (x >= 0 && x < width) intersections.emplace_back(x, height - 1);
69 : }
70 0 : if (vx != 0) {
71 0 : float t = -x0 / vx;
72 0 : float y = y0 + vy * t;
73 0 : if (y >= 0 && y < height) intersections.emplace_back(0, y);
74 0 : t = (width - 1 - x0) / vx;
75 0 : y = y0 + vy * t;
76 0 : if (y >= 0 && y < height) intersections.emplace_back(width - 1, y);
77 : }
78 :
79 0 : if (intersections.size() >= 2) {
80 0 : lanes.push_back({ static_cast<int>(lanes.size() + 1), angle, centroid, {intersections[0], intersections[1]}, cnt });
81 : }
82 : }
83 : }
84 :
85 0 : return lanes;
86 : }
87 :
88 0 : bool LanePostProcessor::shouldMerge(const LaneInfo& a, const LaneInfo& b) const {
89 0 : float angleDiff = std::fabs(a.angle - b.angle);
90 0 : angleDiff = std::min(angleDiff, 180.0f - angleDiff);
91 :
92 0 : if (angleDiff < angleThreshold) {
93 0 : for (const auto& pt1 : a.targets) {
94 0 : for (const auto& pt2 : b.targets) {
95 0 : if (cv::norm(pt1 - pt2) < mergeDistance)
96 0 : return true;
97 : }
98 : }
99 : }
100 0 : return false;
101 : }
102 :
103 0 : std::pair<cv::Point, cv::Point> LanePostProcessor::getExtremities(const std::vector<cv::Point>& contour) const {
104 0 : double maxDist = 0;
105 0 : cv::Point ext1 = contour[0], ext2 = contour[0];
106 :
107 0 : for (size_t i = 0; i < contour.size(); ++i) {
108 0 : for (size_t j = i + 1; j < contour.size(); ++j) {
109 0 : double dist = cv::norm(contour[i] - contour[j]);
110 0 : if (dist > maxDist) {
111 0 : maxDist = dist;
112 0 : ext1 = contour[i];
113 0 : ext2 = contour[j];
114 : }
115 : }
116 : }
117 :
118 0 : return {ext1, ext2};
119 : }
120 :
121 0 : void LanePostProcessor::drawLaneConnections(const std::vector<LaneInfo>& lanes, cv::Mat& canvas) const {
122 0 : for (size_t i = 0; i < lanes.size(); ++i) {
123 0 : for (size_t j = i + 1; j < lanes.size(); ++j) {
124 0 : if (shouldMerge(lanes[i], lanes[j])) {
125 0 : auto [a1, a2] = getExtremities(lanes[i].contour);
126 0 : auto [b1, b2] = getExtremities(lanes[j].contour);
127 :
128 : std::vector<std::pair<cv::Point, cv::Point>> pairs = {
129 : {a1, b1}, {a1, b2}, {a2, b1}, {a2, b2}
130 0 : };
131 :
132 0 : auto bestPair = *std::min_element(pairs.begin(), pairs.end(), [](const auto& p1, const auto& p2) {
133 0 : return cv::norm(p1.first - p1.second) < cv::norm(p2.first - p2.second);
134 0 : });
135 :
136 0 : cv::line(canvas, bestPair.first, bestPair.second, cv::Scalar(255, 255, 255), 3);
137 : }
138 : }
139 : }
140 0 : }
141 :
142 0 : cv::Mat LanePostProcessor::renderFilteredMask(const std::vector<LaneInfo>& lanes, cv::Size shape) const {
143 0 : cv::Mat canvas(shape, CV_8UC3, cv::Scalar(0, 0, 0));
144 0 : for (const auto& lane : lanes) {
145 0 : cv::drawContours(canvas, std::vector<std::vector<cv::Point>>{lane.contour}, -1, cv::Scalar(255, 255, 255), cv::FILLED);
146 : }
147 0 : return canvas;
148 : }
|
