#include #include #include #include #include #include #include "common.hpp" using namespace cv; using namespace std; using namespace dnn; 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 classes; vector colors; void showLegend(); void colorizeSegmentation(const Mat &score, Mat &segm); int main(int argc, char **argv) { CommandLineParser parser(argc, argv, keys); const string modelName = parser.get("@alias"); const string zooFile = parser.get("zoo"); keys += genPreprocArguments(modelName, zooFile); parser = CommandLineParser(argc, argv, keys); parser.about("Use this script to run semantic segmentation deep learning networks using OpenCV."); if (argc == 1 || parser.has("help")) { parser.printMessage(); return 0; } float scale = parser.get("scale"); Scalar mean = parser.get("mean"); bool swapRB = parser.get("rgb"); int inpWidth = parser.get("width"); int inpHeight = parser.get("height"); String model = findFile(parser.get("model")); int backendId = parser.get("backend"); int targetId = parser.get("target"); // Open file with classes names. if (parser.has("classes")) { string file = parser.get("classes"); ifstream ifs(file.c_str()); if (!ifs.is_open()) CV_Error(Error::StsError, "File " + file + " not found"); string line; while (getline(ifs, line)) { classes.push_back(line); } } // Open file with colors. if (parser.has("colors")) { string file = parser.get("colors"); ifstream ifs(file.c_str()); if (!ifs.is_open()) CV_Error(Error::StsError, "File " + file + " not found"); string line; while (getline(ifs, line)) { istringstream colorStr(line.c_str()); Vec3b color; for (int i = 0; i < 3 && !colorStr.eof(); ++i) colorStr >> color[i]; colors.push_back(color); } } if (!parser.check()) { parser.printErrors(); return 1; } CV_Assert(!model.empty()); //! [Read and initialize network] Net net = readNetFromONNX(model); net.setPreferableBackend(backendId); net.setPreferableTarget(targetId); //! [Read and initialize network] // Create a window 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(findFile(parser.get("input"))); else cap.open(parser.get("device")); //! [Open a video file or an image file or a camera stream] // Process frames. Mat frame, blob; while (waitKey(1) < 0) { cap >> frame; if (frame.empty()) { waitKey(); break; } imshow("Original Image", frame); //! [Create a 4D blob from a frame] blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, false); //! [Set input blob] net.setInput(blob); //! [Make forward pass] Mat score = net.forward(); 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(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. vector layersTimes; double freq = getTickFrequency() / 1000; double t = net.getPerfProfile(layersTimes) / freq; 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); if (!classes.empty()) showLegend(); } return 0; } void colorizeSegmentation(const Mat &score, Mat &segm) { const int rows = score.size[2]; const int cols = score.size[3]; const int chns = score.size[1]; if (colors.empty()) { // Generate colors. colors.push_back(Vec3b()); for (int i = 1; i < chns; ++i) { Vec3b color; for (int j = 0; j < 3; ++j) color[j] = (colors[i - 1][j] + rand() % 256) / 2; colors.push_back(color); } } 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())); } Mat maxCl = Mat::zeros(rows, cols, CV_8UC1); Mat maxVal(rows, cols, CV_32FC1, score.data); for (int ch = 1; ch < chns; ch++) { for (int row = 0; row < rows; row++) { const float *ptrScore = score.ptr(0, ch, row); uint8_t *ptrMaxCl = maxCl.ptr(row); float *ptrMaxVal = maxVal.ptr(row); for (int col = 0; col < cols; col++) { if (ptrScore[col] > ptrMaxVal[col]) { ptrMaxVal[col] = ptrScore[col]; ptrMaxCl[col] = (uchar)ch; } } } } segm.create(rows, cols, CV_8UC3); for (int row = 0; row < rows; row++) { const uchar *ptrMaxCl = maxCl.ptr(row); Vec3b *ptrSegm = segm.ptr(row); for (int col = 0; col < cols; col++) { ptrSegm[col] = colors[ptrMaxCl[col]]; } } } void showLegend() { static const int kBlockHeight = 30; static Mat legend; if (legend.empty()) { const int numClasses = (int)classes.size(); 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())); } legend.create(kBlockHeight * numClasses, 200, CV_8UC3); for (int i = 0; i < numClasses; i++) { Mat block = legend.rowRange(i * kBlockHeight, (i + 1) * kBlockHeight); block.setTo(colors[i]); putText(block, classes[i], Point(0, kBlockHeight / 2), FONT_HERSHEY_SIMPLEX, 0.5, Vec3b(255, 255, 255)); } namedWindow("Legend", WINDOW_NORMAL); imshow("Legend", legend); } }