|
1 | 1 | // |
2 | | -// Copyright (c) 2020-2021, Arm Limited and Contributors. All rights reserved. |
| 2 | +// Copyright (c) 2020-2022, Arm Limited and Contributors. All rights reserved. |
3 | 3 | // |
4 | 4 | // SPDX-License-Identifier: Apache-2.0 |
5 | 5 | // |
|
60 | 60 | #define DISPLAY_WINDOW_NAME "YOLO Output" |
61 | 61 |
|
62 | 62 | int main(int argc, char const *argv[]) { |
| 63 | + bool save_as_image = false; |
| 64 | + |
| 65 | + if (argc == 2) { |
| 66 | + save_as_image = true; |
| 67 | + } else if (argc > 2) { |
| 68 | + std::cerr << "Too many arguments have been provided. Please use " |
| 69 | + << "only one argument to save the result of the pipeline as " |
| 70 | + << "an image, or none to display it to an X11 window." << std::endl; |
| 71 | + return 1; |
| 72 | + } |
| 73 | + |
63 | 74 | // load compiled functions |
64 | 75 | tvm::runtime::Module mod = |
65 | 76 | tvm::runtime::Module::LoadFromFile(network_module_path); |
@@ -289,12 +300,21 @@ int main(int argc, char const *argv[]) { |
289 | 300 | } |
290 | 301 | } |
291 | 302 |
|
292 | | - // show in a pop up window the detection results |
293 | | - cv::namedWindow(DISPLAY_WINDOW_NAME, cv::WINDOW_AUTOSIZE); |
294 | | - cv::imshow(DISPLAY_WINDOW_NAME, image); |
| 303 | + if (save_as_image) { |
| 304 | + // save the detection results as an image |
| 305 | + try { |
| 306 | + cv::imwrite(argv[1], image); |
| 307 | + } catch(cv::Exception& e) { |
| 308 | + std::cerr << "An error has occurred while saving to file: " << e.err << std::endl; |
| 309 | + } |
| 310 | + } else { |
| 311 | + // show in a pop up window the detection results |
| 312 | + cv::namedWindow(DISPLAY_WINDOW_NAME, cv::WINDOW_AUTOSIZE); |
| 313 | + cv::imshow(DISPLAY_WINDOW_NAME, image); |
295 | 314 |
|
296 | | - // wait for user to close the window |
297 | | - cv::waitKey(0); |
| 315 | + // wait for user to close the window |
| 316 | + cv::waitKey(0); |
| 317 | + } |
298 | 318 |
|
299 | 319 | // usually the detection results would be filtered again by a non-maximum |
300 | 320 | // supression algorithm. It is omitted here for simplicity. |
|
0 commit comments