1
0
mirror of https://github.com/opencv/opencv.git synced 2026-07-30 07:43:03 +04:00

Merge pull request #25756 from gursimarsingh:bug_fix/segmentation_sample

[BUG FIX] Segmentation sample u2netp model results #25756

PR resloves #25753 related to incorrect output from u2netp model in segmentation sample

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Gursimar Singh
2024-07-03 16:33:12 +05:30
committed by GitHub
parent 55a2a945b6
commit 96a8e6d76c
3 changed files with 34 additions and 31 deletions
+23 -18
View File
@@ -79,7 +79,7 @@ int main(int argc, char **argv)
// Open file with classes names.
if (parser.has("classes"))
{
string file = parser.get<String>("classes");
string file = findFile(parser.get<String>("classes"));
ifstream ifs(file.c_str());
if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found");
@@ -92,7 +92,7 @@ int main(int argc, char **argv)
// Open file with colors.
if (parser.has("colors"))
{
string file = parser.get<String>("colors");
string file = findFile(parser.get<String>("colors"));
ifstream ifs(file.c_str());
if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found");
@@ -146,29 +146,34 @@ int main(int argc, char **argv)
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false);
//! [Set input blob]
net.setInput(blob);
//! [Make forward pass]
Mat score = net.forward();
//! [Set input blob]
if (modelName == "u2netp")
{
Mat mask, thresholded_mask, foreground_overlay, background_overlay, foreground_segmented;
mask = cv::Mat(score.size[2], score.size[3], CV_32F, score.ptr<float>(0, 0));
mask.convertTo(mask, CV_8U, 255);
threshold(mask, thresholded_mask, 0, 255, THRESH_BINARY + THRESH_OTSU);
resize(thresholded_mask, thresholded_mask, Size(frame.cols, frame.rows), 0, 0, INTER_AREA);
vector<Mat> output;
net.forward(output, net.getUnconnectedOutLayersNames());
Mat pred = output[0].reshape(1, output[0].size[2]);
pred.convertTo(pred, CV_8U, 255.0);
Mat mask;
resize(pred, mask, Size(frame.cols, frame.rows), 0, 0, INTER_AREA);
// Create overlays for foreground and background
foreground_overlay = Mat::zeros(frame.size(), frame.type());
background_overlay = Mat::zeros(frame.size(), frame.type());
// Set foreground (object) to red and background to blue
foreground_overlay.setTo(Scalar(0, 0, 255), thresholded_mask);
Mat inverted_mask;
bitwise_not(thresholded_mask, inverted_mask);
background_overlay.setTo(Scalar(255, 0, 0), inverted_mask);
Mat foreground_overlay;
// Set foreground (object) to red
Mat all_zeros = Mat::zeros(frame.size(), CV_8UC1);
vector<Mat> channels = {all_zeros, all_zeros, mask};
merge(channels, foreground_overlay);
// Blend the overlays with the original frame
addWeighted(frame, 1, foreground_overlay, 0.5, 0, foreground_segmented);
addWeighted(foreground_segmented, 1, background_overlay, 0.5, 0, frame);
addWeighted(frame, 0.25, foreground_overlay, 0.75, 0, frame);
}
else
{
//! [Make forward pass]
Mat score = net.forward();
//! [Make forward pass]
Mat segm;
colorizeSegmentation(score, segm);
resize(segm, segm, frame.size(), 0, 0, INTER_NEAREST);