Skip to content

Commit

Permalink
Merge pull request #25559 from gursimarsingh:improved_segmentation_sa…
Browse files Browse the repository at this point in the history 
…mple

Improved segmentation sample #25559

#25006

This pull request replaces caffe models with onnx for the dnn segmentation sample in cpp and python
fcnresnet-50 and fcnresnet-101 has been replaced
u2netp (foreground-background) segmentation onnx model has been added [U2NET](https://github.com/xuebinqin/U-2-Net) 

### Pull Request Readiness Checklist

- [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
- [ ] 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
  • Loading branch information
@gursimarsingh
gursimarsingh committed May 15, 2024
1 parent 17e6b3f commit 48c31bd
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 100 deletions.
25 changes: 18 additions & 7 deletions samples/dnn/models.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,14 +227,13 @@ googlenet:
# Semantic segmentation models.
################################################################################

fcn8s:
fcnresnet50:
load_info:
url: "http://dl.caffe.berkeleyvision.org/fcn8s-heavy-pascal.caffemodel"
sha1: "c449ea74dd7d83751d1357d6a8c323fcf4038962"
model: "fcn8s-heavy-pascal.caffemodel"
config: "fcn8s-heavy-pascal.prototxt"
mean: [0, 0, 0]
scale: 1.0
url: "https://github.com/onnx/models/raw/491ce05590abb7551d7fae43c067c060eeb575a6/validated/vision/object_detection_segmentation/fcn/model/fcn-resnet50-12.onnx"
sha1: "1bb0c7e0034038969aecc6251166f1612a139230"
model: "fcn-resnet50-12.onnx"
mean: [103.5, 116.2, 123.6]
scale: 0.019
width: 500
height: 500
rgb: false
Expand All @@ -251,3 +250,15 @@ fcnresnet101:
height: 500
rgb: false
sample: "segmentation"

u2netp:
load_info:
url: "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
sha1: "0a99236f0d5c1916a99a8c401b23e5ef32038606"
model: "u2netp.onnx"
mean: [123.6, 116.2, 103.5]
scale: 0.019
width: 320
height: 320
rgb: true
sample: "segmentation"
149 changes: 83 additions & 66 deletions samples/dnn/segmentation.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,56 +1,61 @@
#include <fstream>
#include <sstream>
#include <iostream>

#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

#include "common.hpp"

std::string param_keys =
"{ help h | | Print help message. }"
"{ @alias | | An alias name of model to extract preprocessing parameters from models.yml file. }"
"{ zoo | models.yml | An optional path to file with preprocessing parameters }"
"{ device | 0 | camera device number. }"
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
"{ framework f | | Optional name of an origin framework of the model. Detect it automatically if it does not set. }"
"{ classes | | Optional path to a text file with names of classes. }"
"{ colors | | Optional path to a text file with colors for an every class. "
"An every color is represented with three values from 0 to 255 in BGR channels order. }";
std::string backend_keys = cv::format(
"{ backend | 0 | Choose one of computation backends: "
"%d: automatically (by default), "
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
"%d: OpenCV implementation, "
"%d: VKCOM, "
"%d: CUDA }", cv::dnn::DNN_BACKEND_DEFAULT, cv::dnn::DNN_BACKEND_INFERENCE_ENGINE, cv::dnn::DNN_BACKEND_OPENCV, cv::dnn::DNN_BACKEND_VKCOM, cv::dnn::DNN_BACKEND_CUDA);
std::string target_keys = cv::format(
"{ target | 0 | Choose one of target computation devices: "
"%d: CPU target (by default), "
"%d: OpenCL, "
"%d: OpenCL fp16 (half-float precision), "
"%d: VPU, "
"%d: Vulkan, "
"%d: CUDA, "
"%d: CUDA fp16 (half-float preprocess) }", cv::dnn::DNN_TARGET_CPU, cv::dnn::DNN_TARGET_OPENCL, cv::dnn::DNN_TARGET_OPENCL_FP16, cv::dnn::DNN_TARGET_MYRIAD, cv::dnn::DNN_TARGET_VULKAN, cv::dnn::DNN_TARGET_CUDA, cv::dnn::DNN_TARGET_CUDA_FP16);
std::string keys = param_keys + backend_keys + target_keys;

using namespace cv;
using namespace std;
using namespace dnn;

std::vector<std::string> classes;
std::vector<Vec3b> colors;
const string param_keys =
"{ help h | | Print help message. }"
"{ @alias | | An alias name of model to extract preprocessing parameters from models.yml file. }"
"{ zoo | models.yml | An optional path to file with preprocessing parameters }"
"{ device | 0 | camera device number. }"
"{ input i | | Path to input image or video file. Skip this argument to capture frames from a camera. }"
"{ classes | | Optional path to a text file with names of classes. }"
"{ colors | | Optional path to a text file with colors for an every class. "
"Every color is represented with three values from 0 to 255 in BGR channels order. }";

const string backend_keys = format(
"{ backend | 0 | Choose one of computation backends: "
"%d: automatically (by default), "
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
"%d: OpenCV implementation, "
"%d: VKCOM, "
"%d: CUDA }",
DNN_BACKEND_DEFAULT, DNN_BACKEND_INFERENCE_ENGINE, DNN_BACKEND_OPENCV, DNN_BACKEND_VKCOM, DNN_BACKEND_CUDA);

const string target_keys = format(
"{ target | 0 | Choose one of target computation devices: "
"%d: CPU target (by default), "
"%d: OpenCL, "
"%d: OpenCL fp16 (half-float precision), "
"%d: VPU, "
"%d: Vulkan, "
"%d: CUDA, "
"%d: CUDA fp16 (half-float preprocess) }",
DNN_TARGET_CPU, DNN_TARGET_OPENCL, DNN_TARGET_OPENCL_FP16, DNN_TARGET_MYRIAD, DNN_TARGET_VULKAN, DNN_TARGET_CUDA, DNN_TARGET_CUDA_FP16);

string keys = param_keys + backend_keys + target_keys;
vector<string> classes;
vector<Vec3b> colors;

void showLegend();

void colorizeSegmentation(const Mat &score, Mat &segm);

int main(int argc, char** argv)
int main(int argc, char **argv)
{
CommandLineParser parser(argc, argv, keys);

const std::string modelName = parser.get<String>("@alias");
const std::string zooFile = parser.get<String>("zoo");
const string modelName = parser.get<String>("@alias");
const string zooFile = parser.get<String>("zoo");

keys += genPreprocArguments(modelName, zooFile);

Expand All @@ -68,36 +73,33 @@ int main(int argc, char** argv)
int inpWidth = parser.get<int>("width");
int inpHeight = parser.get<int>("height");
String model = findFile(parser.get<String>("model"));
String config = findFile(parser.get<String>("config"));
String framework = parser.get<String>("framework");
int backendId = parser.get<int>("backend");
int targetId = parser.get<int>("target");

// Open file with classes names.
if (parser.has("classes"))
{
std::string file = parser.get<String>("classes");
std::ifstream ifs(file.c_str());
string file = parser.get<String>("classes");
ifstream ifs(file.c_str());
if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found");
std::string line;
while (std::getline(ifs, line))
string line;
while (getline(ifs, line))
{
classes.push_back(line);
}
}

// Open file with colors.
if (parser.has("colors"))
{
std::string file = parser.get<String>("colors");
std::ifstream ifs(file.c_str());
string file = parser.get<String>("colors");
ifstream ifs(file.c_str());
if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found");
std::string line;
while (std::getline(ifs, line))
string line;
while (getline(ifs, line))
{
std::istringstream colorStr(line.c_str());
istringstream colorStr(line.c_str());

Vec3b color;
for (int i = 0; i < 3 && !colorStr.eof(); ++i)
Expand All @@ -114,23 +116,21 @@ int main(int argc, char** argv)

CV_Assert(!model.empty());
//! [Read and initialize network]
Net net = readNet(model, config, framework);
Net net = readNetFromONNX(model);
net.setPreferableBackend(backendId);
net.setPreferableTarget(targetId);
//! [Read and initialize network]

// Create a window
static const std::string kWinName = "Deep learning semantic segmentation in OpenCV";
static const string kWinName = "Deep learning semantic segmentation in OpenCV";
namedWindow(kWinName, WINDOW_NORMAL);

//! [Open a video file or an image file or a camera stream]
VideoCapture cap;
if (parser.has("input"))
cap.open(parser.get<String>("input"));
cap.open(findFile(parser.get<String>("input")));
else
cap.open(parser.get<int>("device"));
//! [Open a video file or an image file or a camera stream]

// Process frames.
Mat frame, blob;
while (waitKey(1) < 0)
Expand All @@ -141,29 +141,45 @@ int main(int argc, char** argv)
waitKey();
break;
}

imshow("Original Image", frame);
//! [Create a 4D blob from a frame]
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false);
//! [Create a 4D blob from a frame]

//! [Set input blob]
net.setInput(blob);
//! [Set input blob]
//! [Make forward pass]
Mat score = net.forward();
//! [Make forward pass]

Mat segm;
colorizeSegmentation(score, segm);

resize(segm, segm, frame.size(), 0, 0, INTER_NEAREST);
addWeighted(frame, 0.1, segm, 0.9, 0.0, frame);
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);
// 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);
// 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);
}
else
{
Mat segm;
colorizeSegmentation(score, segm);
resize(segm, segm, frame.size(), 0, 0, INTER_NEAREST);
addWeighted(frame, 0.1, segm, 0.9, 0.0, frame);
}

// Put efficiency information.
std::vector<double> layersTimes;
vector<double> layersTimes;
double freq = getTickFrequency() / 1000;
double t = net.getPerfProfile(layersTimes) / freq;
std::string label = format("Inference time: %.2f ms", t);
string label = format("Inference time: %.2f ms", t);
putText(frame, label, Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));

imshow(kWinName, frame);
Expand Down Expand Up @@ -194,7 +210,8 @@ void colorizeSegmentation(const Mat &score, Mat &segm)
else if (chns != (int)colors.size())
{
CV_Error(Error::StsError, format("Number of output classes does not match "
"number of colors (%d != %zu)", chns, colors.size()));
"number of colors (%d != %zu)",
chns, colors.size()));
}

Mat maxCl = Mat::zeros(rows, cols, CV_8UC1);
Expand All @@ -216,7 +233,6 @@ void colorizeSegmentation(const Mat &score, Mat &segm)
}
}
}

segm.create(rows, cols, CV_8UC3);
for (int row = 0; row < rows; row++)
{
Expand All @@ -239,7 +255,8 @@ void showLegend()
if ((int)colors.size() != numClasses)
{
CV_Error(Error::StsError, format("Number of output classes does not match "
"number of labels (%zu != %zu)", colors.size(), classes.size()));
"number of labels (%zu != %zu)",
colors.size(), classes.size()));
}
legend.create(kBlockHeight * numClasses, 200, CV_8UC3);
for (int i = 0; i < numClasses; i++)
Expand Down
62 changes: 35 additions & 27 deletions samples/dnn/segmentation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,6 @@
parser.add_argument('--zoo', default=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models.yml'),
help='An optional path to file with preprocessing parameters.')
parser.add_argument('--input', help='Path to input image or video file. Skip this argument to capture frames from a camera.')
parser.add_argument('--framework', choices=['caffe', 'tensorflow', 'darknet', 'onnx'],
help='Optional name of an origin framework of the model. '
'Detect it automatically if it does not set.')
parser.add_argument('--colors', help='Optional path to a text file with colors for an every class. '
'An every color is represented with three values from 0 to 255 in BGR channels order.')
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
Expand Down Expand Up @@ -44,7 +41,6 @@
args = parser.parse_args()

args.model = findFile(args.model)
args.config = findFile(args.config)
args.classes = findFile(args.classes)

np.random.seed(324)
Expand Down Expand Up @@ -79,7 +75,7 @@ def showLegend(classes):
classes = None

# Load a network
net = cv.dnn.readNet(args.model, args.config, args.framework)
net = cv.dnn.readNet(args.model)
net.setPreferableBackend(args.backend)
net.setPreferableTarget(args.target)

Expand All @@ -94,41 +90,53 @@ def showLegend(classes):
cv.waitKey()
break

cv.imshow("Original Image", frame)
frameHeight = frame.shape[0]
frameWidth = frame.shape[1]

# Create a 4D blob from a frame.
inpWidth = args.width if args.width else frameWidth
inpHeight = args.height if args.height else frameHeight
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=False)

# Run a model
net.setInput(blob)
score = net.forward()

numClasses = score.shape[1]
height = score.shape[2]
width = score.shape[3]

# Draw segmentation
if not colors:
# Generate colors
colors = [np.array([0, 0, 0], np.uint8)]
for i in range(1, numClasses):
colors.append((colors[i - 1] + np.random.randint(0, 256, [3], np.uint8)) / 2)

classIds = np.argmax(score[0], axis=0)
segm = np.stack([colors[idx] for idx in classIds.flatten()])
segm = segm.reshape(height, width, 3)

segm = cv.resize(segm, (frameWidth, frameHeight), interpolation=cv.INTER_NEAREST)
frame = (0.1 * frame + 0.9 * segm).astype(np.uint8)
if args.alias == 'u2netp':
mask = score[0][0]
mask = mask.astype(np.uint8)
_, mask = cv.threshold(mask, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
mask = cv.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv.INTER_AREA)
# Create overlays for foreground and background
foreground_overlay = np.zeros_like(frame, dtype=np.uint8)
background_overlay = np.zeros_like(frame, dtype=np.uint8)
# Set foreground (object) to red and background to blue
foreground_overlay[mask == 255] = [0, 0, 255] # Red foreground
background_overlay[mask == 0] = [255, 0, 0] # Blue background
# Blend the overlays with the original frame
foreground_segmented = cv.addWeighted(frame, 1, foreground_overlay, 0.5, 0)
frame = cv.addWeighted(foreground_segmented, 1, background_overlay, 0.5, 0)
else:
numClasses = score.shape[1]
height = score.shape[2]
width = score.shape[3]
# Draw segmentation
if not colors:
# Generate colors
colors = [np.array([0, 0, 0], np.uint8)]
for i in range(1, numClasses):
colors.append((colors[i - 1] + np.random.randint(0, 256, [3], np.uint8)) / 2)
classIds = np.argmax(score[0], axis=0)
segm = np.stack([colors[idx] for idx in classIds.flatten()])
segm = segm.reshape(height, width, 3)

segm = cv.resize(segm, (frameWidth, frameHeight), interpolation=cv.INTER_NEAREST)
frame = (0.1 * frame + 0.9 * segm).astype(np.uint8)

showLegend(classes)

# Put efficiency information.
t, _ = net.getPerfProfile()
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency())
cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))

showLegend(classes)

cv.imshow(winName, frame)
cv.imshow(winName, frame)

0 comments on commit 48c31bd

Please sign in to comment.