samples/dnn/classification.cpp#

Check the corresponding tutorial for more details

  1#include <fstream>
  2#include <sstream>
  3#include <iostream>
  4
  5#include <opencv2/dnn.hpp>
  6#include <opencv2/imgproc.hpp>
  7#include <opencv2/highgui.hpp>
  8#include <opencv2/core/utils/logger.hpp>
  9
 10#include "common.hpp"
 11
 12using namespace cv;
 13using namespace std;
 14using namespace dnn;
 15
 16const string about =
 17        "Use this script to run a classification model on a camera stream, video, image or image list (i.e. .xml or .yaml containing image lists)\n\n"
 18        "Firstly, download required models using `download_models.py` (if not already done). Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to specify where models should be downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.\n"
 19        "To run:\n"
 20        "\t ./example_dnn_classification model_name --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)\n"
 21        "Sample command:\n"
 22        "\t ./example_dnn_classification resnet --input=$OPENCV_SAMPLES_DATA_PATH/baboon.jpg\n"
 23        "\t ./example_dnn_classification squeezenet\n"
 24        "Model path can also be specified using --model argument. "
 25        "Use imagelist_creator to create the xml or yaml list\n";
 26
 27const string param_keys =
 28    "{ help  h         |                   | Print help message. }"
 29    "{ @alias          |                   | An alias name of model to extract preprocessing parameters from models.yml file. }"
 30    "{ zoo             | ../dnn/models.yml | An optional path to file with preprocessing parameters }"
 31    "{ input i         |                   | Path to input image or video file. Skip this argument to capture frames from a camera.}"
 32    "{ imglist         |                   | Pass this flag if image list (i.e. .xml or .yaml) file is passed}"
 33    "{ crop            |       false       | Preprocess input image by center cropping.}"
 34    //"{ labels          |                   | Path to the text file with labels for detected objects.}"
 35    "{ model           |                   | Path to the model file.}";
 36
 37const string backend_keys = format(
 38    "{ backend          | default | Choose one of computation backends: "
 39                              "default: automatically (by default), "
 40                              "openvino: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
 41                              "opencv: OpenCV implementation, "
 42                              "vkcom: VKCOM, "
 43                              "cuda: CUDA, "
 44                              "webnn: WebNN }");
 45
 46const string target_keys = format(
 47    "{ target           | cpu | Choose one of target computation devices: "
 48                              "cpu: CPU target (by default), "
 49                              "opencl: OpenCL, "
 50                              "opencl_fp16: OpenCL fp16 (half-float precision), "
 51                              "vpu: VPU, "
 52                              "vulkan: Vulkan, "
 53                              "cuda: CUDA, "
 54                              "cuda_fp16: CUDA fp16 (half-float preprocess) }");
 55
 56string keys = param_keys + backend_keys + target_keys;
 57
 58vector<string> classes;
 59static bool readStringList( const string& filename, vector<string>& l )
 60{
 61    l.resize(0);
 62    FileStorage fs(filename, FileStorage::READ);
 63    if( !fs.isOpened() )
 64        return false;
 65    size_t dir_pos = filename.rfind('/');
 66    if (dir_pos == string::npos)
 67        dir_pos = filename.rfind('\\');
 68    FileNode n = fs.getFirstTopLevelNode();
 69    if( n.type() != FileNode::SEQ )
 70        return false;
 71    FileNodeIterator it = n.begin(), it_end = n.end();
 72    for( ; it != it_end; ++it )
 73    {
 74        string fname = (string)*it;
 75        if (dir_pos != string::npos)
 76        {
 77            string fpath = samples::findFile(filename.substr(0, dir_pos + 1) + fname, false);
 78            if (fpath.empty())
 79            {
 80                fpath = samples::findFile(fname);
 81            }
 82            fname = fpath;
 83        }
 84        else
 85        {
 86            fname = samples::findFile(fname);
 87        }
 88        l.push_back(fname);
 89    }
 90    return true;
 91}
 92
 93int main(int argc, char** argv)
 94{
 95    utils::logging::setLogLevel(utils::logging::LOG_LEVEL_INFO);
 96
 97    CommandLineParser parser(argc, argv, keys);
 98
 99    if (!parser.has("@alias") || parser.has("help"))
100    {
101        cout << about << endl;
102        parser.printMessage();
103        return -1;
104    }
105    const string modelName = parser.get<String>("@alias");
106    const string zooFile = findFile(parser.get<String>("zoo"));
107
108    keys += genPreprocArguments(modelName, zooFile);
109    parser = CommandLineParser(argc, argv, keys);
110    parser.about(about);
111    if (argc == 1 || parser.has("help"))
112    {
113        parser.printMessage();
114        return 0;
115    }
116    String sha1 = parser.get<String>("sha1");
117    float scale = parser.get<float>("scale");
118    Scalar mean = parser.get<Scalar>("mean");
119    Scalar std = parser.get<Scalar>("std");
120    bool swapRB = parser.get<bool>("rgb");
121    bool crop = parser.get<bool>("crop");
122    int inpWidth = parser.get<int>("width");
123    int inpHeight = parser.get<int>("height");
124    String model = findModel(parser.get<String>("model"), sha1);
125    String backend = parser.get<String>("backend");
126    String target = parser.get<String>("target");
127    bool isImgList = parser.has("imglist");
128
129    // Open file with labels.
130    string labels_filename = parser.get<String>("labels");
131    string file = findFile(labels_filename);
132    ifstream ifs(file.c_str());
133    if (!ifs.is_open()){
134        cout<<"File " << file << " not found";
135        exit(1);
136    }
137    string line;
138    while (getline(ifs, line))
139    {
140        classes.push_back(line);
141    }
142    if (!parser.check())
143    {
144        parser.printErrors();
145        return 1;
146    }
147    CV_Assert(!model.empty());
148    //! [Read and initialize network]
149    EngineType engine = ENGINE_AUTO;
150    if (backend != "default" || target != "cpu"){
151        engine = ENGINE_CLASSIC;
152    }
153    Net net = readNetFromONNX(model, engine);
154    net.setPreferableBackend(getBackendID(backend));
155    net.setPreferableTarget(getTargetID(target));
156    net.setProfilingMode(DNN_PROFILE_SUMMARY);
157    //! [Read and initialize network]
158
159    // Create a window
160    static const std::string kWinName = "Deep learning image classification in OpenCV";
161    namedWindow(kWinName, WINDOW_NORMAL);
162
163    //Create FontFace for putText
164    FontFace sans("sans");
165
166    //! [Open a video file or an image file or a camera stream]
167    VideoCapture cap;
168    vector<string> imageList;
169    size_t currentImageIndex = 0;
170
171    if (parser.has("input")) {
172        string input = findFile(parser.get<String>("input"));
173
174        if (isImgList) {
175            bool check = readStringList(samples::findFile(input), imageList);
176            if (imageList.empty() || !check) {
177                cout << "Error: No images found or the provided file is not a valid .yaml or .xml file." << endl;
178                return -1;
179            }
180        } else {
181            // Input is not a directory, try to open as video or image
182            cap.open(input);
183            if (!cap.isOpened()) {
184                cout << "Failed to open the input." << endl;
185                return -1;
186            }
187        }
188    } else {
189        cap.open(0); // Open default camera
190    }
191    //! [Open a video file or an image file or a camera stream]
192
193    Mat frame, blob;
194    for(;;)
195    {
196        if (!imageList.empty()) {
197            // Handling directory of images
198            if (currentImageIndex >= imageList.size()) {
199                waitKey();
200                break; // Exit if all images are processed
201            }
202            frame = imread(imageList[currentImageIndex++]);
203            if(frame.empty()){
204                cout<<"Cannot open file"<<endl;
205                continue;
206            }
207        } else {
208            // Handling video or single image
209            cap >> frame;
210        }
211        if (frame.empty())
212        {
213            break;
214        }
215        //! [Create a 4D blob from a frame]
216        blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, crop);
217        // Check std values.
218        if (std.val[0] != 0.0 && std.val[1] != 0.0 && std.val[2] != 0.0)
219        {
220            // Divide blob by std.
221            divide(blob, std, blob);
222        }
223        //! [Create a 4D blob from a frame]
224        //! [Set input blob]
225        net.setInput(blob);
226        //! [Set input blob]
227
228        TickMeter timeRecorder;
229        timeRecorder.reset();
230        Mat prob = net.forward();
231        double t1;
232        //! [Make forward pass]
233        timeRecorder.start();
234        prob = net.forward();
235        timeRecorder.stop();
236        net.printPerfProfile();
237        //! [Make forward pass]
238
239        //! [Get a class with a highest score]
240        int N = (int)prob.total(), K = std::min(5, N);
241        std::vector<std::pair<float, int> > prob_vec;
242        for (int i = 0; i < N; i++) {
243            prob_vec.push_back(std::make_pair(-prob.at<float>(i), i));
244        }
245        std::sort(prob_vec.begin(), prob_vec.end());
246
247        //! [Get a class with a highest score]
248        t1 = timeRecorder.getTimeMilli();
249        timeRecorder.reset();
250        string label = format("Inference time: %.1f ms", t1);
251        Mat subframe = frame(Rect(0, 0, std::min(1000, frame.cols), std::min(300, frame.rows)));
252        subframe *= 0.3f;
253        putText(frame, label, Point(20, 50), Scalar(0, 255, 0), sans, 25, 800);
254
255        // Print predicted class.
256        for (int i = 0; i < K; i++) {
257            int classId = prob_vec[i].second;
258            float confidence = -prob_vec[i].first;
259            label = format("%d. %s: %.2f", i+1, (classes.empty() ? format("Class #%d", classId).c_str() :
260                                        classes[classId].c_str()), confidence);
261            putText(frame, label, Point(20, 110 + i*35), Scalar(0, 255, 0), sans, 25, 500);
262        }
263        imshow(kWinName, frame);
264        int key = waitKey(isImgList ? 1000 : 100);
265        if (key == ' ')
266            key = waitKey();
267        if (key == 'q' || key == 27) // Check if 'q' or 'ESC' is pressed
268            return 0;
269    }
270    waitKey();
271    return 0;
272}