Chiebot-Mirror
/
opencv
mirror of https://github.com/opencv/opencv.git
8
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #25559 from gursimarsingh:improved_segmentation_sample

Browse Source 
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
pull/25610/head
Gursimar Singh 7 months ago committed by GitHub
parent 17e6b3f931
commit 48c31bddc4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 136 additions and 100 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 25
      samples/dnn/models.yml
  2. 149
      samples/dnn/segmentation.cpp
  3. 62
      samples/dnn/segmentation.py

25
samples/dnn/models.yml
Unescape View File

@ -227,14 +227,13 @@ googlenet:
# Semantic segmentation models. # Semantic segmentation models.
################################################################################ ################################################################################
fcn8s: fcnresnet50:
load_info: load_info:
url: "http://dl.caffe.berkeleyvision.org/fcn8s-heavy-pascal.caffemodel" url: "https://github.com/onnx/models/raw/491ce05590abb7551d7fae43c067c060eeb575a6/validated/vision/object_detection_segmentation/fcn/model/fcn-resnet50-12.onnx"
sha1: "c449ea74dd7d83751d1357d6a8c323fcf4038962" sha1: "1bb0c7e0034038969aecc6251166f1612a139230"
model: "fcn8s-heavy-pascal.caffemodel" model: "fcn-resnet50-12.onnx"
config: "fcn8s-heavy-pascal.prototxt" mean: [103.5, 116.2, 123.6]
mean: [0, 0, 0] scale: 0.019
scale: 1.0
width: 500 width: 500
height: 500 height: 500
rgb: false rgb: false
@ -251,3 +250,15 @@ fcnresnet101:
height: 500 height: 500
rgb: false rgb: false
sample: "segmentation" 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
samples/dnn/segmentation.cpp
Unescape View File

@ -1,5 +1,6 @@
#include <fstream> #include <fstream>
#include <sstream> #include <sstream>
#include <iostream>
#include <opencv2/dnn.hpp> #include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp> #include <opencv2/imgproc.hpp>
@ -7,50 +8,54 @@
#include "common.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 cv;
using namespace std;
using namespace dnn; using namespace dnn;
std::vector<std::string> classes; const string param_keys =
std::vector<Vec3b> colors; "{ 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 showLegend();
void colorizeSegmentation(const Mat &score, Mat &segm); void colorizeSegmentation(const Mat &score, Mat &segm);
int main(int argc, char** argv) int main(int argc, char **argv)
{ {
CommandLineParser parser(argc, argv, keys); CommandLineParser parser(argc, argv, keys);
const std::string modelName = parser.get<String>("@alias"); const string modelName = parser.get<String>("@alias");
const std::string zooFile = parser.get<String>("zoo"); const string zooFile = parser.get<String>("zoo");
keys += genPreprocArguments(modelName, zooFile); keys += genPreprocArguments(modelName, zooFile);
@ -68,36 +73,33 @@ int main(int argc, char** argv)
int inpWidth = parser.get<int>("width"); int inpWidth = parser.get<int>("width");
int inpHeight = parser.get<int>("height"); int inpHeight = parser.get<int>("height");
String model = findFile(parser.get<String>("model")); 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 backendId = parser.get<int>("backend");
int targetId = parser.get<int>("target"); int targetId = parser.get<int>("target");
// Open file with classes names. // Open file with classes names.
if (parser.has("classes")) if (parser.has("classes"))
{ {
std::string file = parser.get<String>("classes"); string file = parser.get<String>("classes");
std::ifstream ifs(file.c_str()); ifstream ifs(file.c_str());
if (!ifs.is_open()) if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found"); CV_Error(Error::StsError, "File " + file + " not found");
std::string line; string line;
while (std::getline(ifs, line)) while (getline(ifs, line))
{ {
classes.push_back(line); classes.push_back(line);
} }
} }
// Open file with colors. // Open file with colors.
if (parser.has("colors")) if (parser.has("colors"))
{ {
std::string file = parser.get<String>("colors"); string file = parser.get<String>("colors");
std::ifstream ifs(file.c_str()); ifstream ifs(file.c_str());
if (!ifs.is_open()) if (!ifs.is_open())
CV_Error(Error::StsError, "File " + file + " not found"); CV_Error(Error::StsError, "File " + file + " not found");
std::string line; string line;
while (std::getline(ifs, line)) while (getline(ifs, line))
{ {
std::istringstream colorStr(line.c_str()); istringstream colorStr(line.c_str());
Vec3b color; Vec3b color;
for (int i = 0; i < 3 && !colorStr.eof(); ++i) for (int i = 0; i < 3 && !colorStr.eof(); ++i)
@ -114,23 +116,21 @@ int main(int argc, char** argv)
CV_Assert(!model.empty()); CV_Assert(!model.empty());
//! [Read and initialize network] //! [Read and initialize network]
Net net = readNet(model, config, framework); Net net = readNetFromONNX(model);
net.setPreferableBackend(backendId); net.setPreferableBackend(backendId);
net.setPreferableTarget(targetId); net.setPreferableTarget(targetId);
//! [Read and initialize network] //! [Read and initialize network]
// Create a window // 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); namedWindow(kWinName, WINDOW_NORMAL);
//! [Open a video file or an image file or a camera stream] //! [Open a video file or an image file or a camera stream]
VideoCapture cap; VideoCapture cap;
if (parser.has("input")) if (parser.has("input"))
cap.open(parser.get<String>("input")); cap.open(findFile(parser.get<String>("input")));
else else
cap.open(parser.get<int>("device")); cap.open(parser.get<int>("device"));
//! [Open a video file or an image file or a camera stream] //! [Open a video file or an image file or a camera stream]
// Process frames. // Process frames.
Mat frame, blob; Mat frame, blob;
while (waitKey(1) < 0) while (waitKey(1) < 0)
@ -141,29 +141,45 @@ int main(int argc, char** argv)
waitKey(); waitKey();
break; break;
} }
imshow("Original Image", frame);
//! [Create a 4D blob from a frame] //! [Create a 4D blob from a frame]
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false); blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false);
//! [Create a 4D blob from a frame]
//! [Set input blob] //! [Set input blob]
net.setInput(blob); net.setInput(blob);
//! [Set input blob]
//! [Make forward pass] //! [Make forward pass]
Mat score = net.forward(); Mat score = net.forward();
//! [Make forward pass] if (modelName == "u2netp")
{
Mat segm; Mat mask, thresholded_mask, foreground_overlay, background_overlay, foreground_segmented;
colorizeSegmentation(score, segm); mask = cv::Mat(score.size[2], score.size[3], CV_32F, score.ptr<float>(0, 0));
mask.convertTo(mask, CV_8U, 255);
resize(segm, segm, frame.size(), 0, 0, INTER_NEAREST); threshold(mask, thresholded_mask, 0, 255, THRESH_BINARY + THRESH_OTSU);
addWeighted(frame, 0.1, segm, 0.9, 0.0, frame); 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. // Put efficiency information.
std::vector<double> layersTimes; vector<double> layersTimes;
double freq = getTickFrequency() / 1000; double freq = getTickFrequency() / 1000;
double t = net.getPerfProfile(layersTimes) / freq; 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)); putText(frame, label, Point(0, 15), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 255, 0));
imshow(kWinName, frame); imshow(kWinName, frame);
@ -194,7 +210,8 @@ void colorizeSegmentation(const Mat &score, Mat &segm)
else if (chns != (int)colors.size()) else if (chns != (int)colors.size())
{ {
CV_Error(Error::StsError, format("Number of output classes does not match " 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); Mat maxCl = Mat::zeros(rows, cols, CV_8UC1);
@ -216,7 +233,6 @@ void colorizeSegmentation(const Mat &score, Mat &segm)
} }
} }
} }
segm.create(rows, cols, CV_8UC3); segm.create(rows, cols, CV_8UC3);
for (int row = 0; row < rows; row++) for (int row = 0; row < rows; row++)
{ {
@ -239,7 +255,8 @@ void showLegend()
if ((int)colors.size() != numClasses) if ((int)colors.size() != numClasses)
{ {
CV_Error(Error::StsError, format("Number of output classes does not match " 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); legend.create(kBlockHeight * numClasses, 200, CV_8UC3);
for (int i = 0; i < numClasses; i++) for (int i = 0; i < numClasses; i++)

62
samples/dnn/segmentation.py
Unescape View File

@ -14,9 +14,6 @@ parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--zoo', default=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models.yml'), 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.') 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('--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. ' 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.') '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, parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
@ -44,7 +41,6 @@ parser = argparse.ArgumentParser(parents=[parser],
args = parser.parse_args() args = parser.parse_args()
args.model = findFile(args.model) args.model = findFile(args.model)
args.config = findFile(args.config)
args.classes = findFile(args.classes) args.classes = findFile(args.classes)
np.random.seed(324) np.random.seed(324)
@ -79,7 +75,7 @@ def showLegend(classes):
classes = None classes = None
# Load a network # Load a network
net = cv.dnn.readNet(args.model, args.config, args.framework) net = cv.dnn.readNet(args.model)
net.setPreferableBackend(args.backend) net.setPreferableBackend(args.backend)
net.setPreferableTarget(args.target) net.setPreferableTarget(args.target)
@ -94,41 +90,53 @@ while cv.waitKey(1) < 0:
cv.waitKey() cv.waitKey()
break break
cv.imshow("Original Image", frame)
frameHeight = frame.shape[0] frameHeight = frame.shape[0]
frameWidth = frame.shape[1] frameWidth = frame.shape[1]
# Create a 4D blob from a frame. # Create a 4D blob from a frame.
inpWidth = args.width if args.width else frameWidth inpWidth = args.width if args.width else frameWidth
inpHeight = args.height if args.height else frameHeight inpHeight = args.height if args.height else frameHeight
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=False) blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=False)
# Run a model
net.setInput(blob) net.setInput(blob)
score = net.forward() score = net.forward()
numClasses = score.shape[1] if args.alias == 'u2netp':
height = score.shape[2] mask = score[0][0]
width = score.shape[3] mask = mask.astype(np.uint8)
_, mask = cv.threshold(mask, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
# Draw segmentation mask = cv.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv.INTER_AREA)
if not colors: # Create overlays for foreground and background
# Generate colors foreground_overlay = np.zeros_like(frame, dtype=np.uint8)
colors = [np.array([0, 0, 0], np.uint8)] background_overlay = np.zeros_like(frame, dtype=np.uint8)
for i in range(1, numClasses): # Set foreground (object) to red and background to blue
colors.append((colors[i - 1] + np.random.randint(0, 256, [3], np.uint8)) / 2) foreground_overlay[mask == 255] = [0, 0, 255] # Red foreground
background_overlay[mask == 0] = [255, 0, 0] # Blue background
classIds = np.argmax(score[0], axis=0) # Blend the overlays with the original frame
segm = np.stack([colors[idx] for idx in classIds.flatten()]) foreground_segmented = cv.addWeighted(frame, 1, foreground_overlay, 0.5, 0)
segm = segm.reshape(height, width, 3) frame = cv.addWeighted(foreground_segmented, 1, background_overlay, 0.5, 0)
else:
segm = cv.resize(segm, (frameWidth, frameHeight), interpolation=cv.INTER_NEAREST) numClasses = score.shape[1]
frame = (0.1 * frame + 0.9 * segm).astype(np.uint8) 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. # Put efficiency information.
t, _ = net.getPerfProfile() t, _ = net.getPerfProfile()
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency()) 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)) 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)
Write Preview
Loading…
Cancel
Save