1#include <iostream> 2#include <string> 3#include <sstream> 4#include <iomanip> 5#include <stdexcept> 6#include <opencv2/core/utility.hpp> 7#include "opencv2/cudastereo.hpp" 8#include "opencv2/highgui.hpp" 9#include "opencv2/imgproc.hpp" 10 11using namespace cv; 12using namespace std; 13 14bool help_showed = false; 15 16struct Params 17{ 18 Params(); 19 static Params read(int argc, char** argv); 20 21 string left; 22 string right; 23 24 string method_str() const 25 { 26 switch (method) 27 { 28 case BM: return "BM"; 29 case BP: return "BP"; 30 case CSBP: return "CSBP"; 31 } 32 return ""; 33 } 34 enum {BM, BP, CSBP} method; 35 int ndisp; // Max disparity + 1 36}; 37 38 39struct App 40{ 41 App(const Params& p); 42 void run(); 43 void handleKey(char key); 44 void printParams() const; 45 46 void workBegin() { work_begin = getTickCount(); } 47 void workEnd() 48 { 49 int64 d = getTickCount() - work_begin; 50 double f = getTickFrequency(); 51 work_fps = f / d; 52 } 53 54 string text() const 55 { 56 stringstream ss; 57 ss << "(" << p.method_str() << ") FPS: " << setiosflags(ios::left) 58 << setprecision(4) << work_fps; 59 return ss.str(); 60 } 61private: 62 Params p; 63 bool running; 64 65 Mat left_src, right_src; 66 Mat left, right; 67 cuda::GpuMat d_left, d_right; 68 69 Ptr<cuda::StereoBM> bm; 70 Ptr<cuda::StereoBeliefPropagation> bp; 71 Ptr<cuda::StereoConstantSpaceBP> csbp; 72 73 int64 work_begin; 74 double work_fps; 75}; 76 77static void printHelp() 78{ 79 cout << "Usage: stereo_match_gpu\n" 80 << "\t--left <left_view> --right <right_view> # must be rectified\n" 81 << "\t--method <stereo_match_method> # BM | BP | CSBP\n" 82 << "\t--ndisp <number> # number of disparity levels\n"; 83 help_showed = true; 84} 85 86int main(int argc, char** argv) 87{ 88 try 89 { 90 if (argc < 2) 91 { 92 printHelp(); 93 return 1; 94 } 95 Params args = Params::read(argc, argv); 96 if (help_showed) 97 return -1; 98 App app(args); 99 app.run(); 100 } 101 catch (const exception& e) 102 { 103 cout << "error: " << e.what() << endl; 104 } 105 return 0; 106} 107 108 109Params::Params() 110{ 111 method = BM; 112 ndisp = 64; 113} 114 115 116Params Params::read(int argc, char** argv) 117{ 118 Params p; 119 120 for (int i = 1; i < argc; i++) 121 { 122 if (string(argv[i]) == "--left") p.left = argv[++i]; 123 else if (string(argv[i]) == "--right") p.right = argv[++i]; 124 else if (string(argv[i]) == "--method") 125 { 126 if (string(argv[i + 1]) == "BM") p.method = BM; 127 else if (string(argv[i + 1]) == "BP") p.method = BP; 128 else if (string(argv[i + 1]) == "CSBP") p.method = CSBP; 129 else throw runtime_error("unknown stereo match method: " + string(argv[i + 1])); 130 i++; 131 } 132 else if (string(argv[i]) == "--ndisp") p.ndisp = atoi(argv[++i]); 133 else if (string(argv[i]) == "--help") printHelp(); 134 else throw runtime_error("unknown key: " + string(argv[i])); 135 } 136 137 return p; 138} 139 140 141App::App(const Params& params) 142 : p(params), running(false) 143{ 144 cv::cuda::printShortCudaDeviceInfo(cv::cuda::getDevice()); 145 146 cout << "stereo_match_gpu sample\n"; 147 cout << "\nControls:\n" 148 << "\tesc - exit\n" 149 << "\tp - print current parameters\n" 150 << "\tg - convert source images into gray\n" 151 << "\tm - change stereo match method\n" 152 << "\ts - change Sobel prefiltering flag (for BM only)\n" 153 << "\t1/q - increase/decrease maximum disparity\n" 154 << "\t2/w - increase/decrease window size (for BM only)\n" 155 << "\t3/e - increase/decrease iteration count (for BP and CSBP only)\n" 156 << "\t4/r - increase/decrease level count (for BP and CSBP only)\n"; 157} 158 159 160void App::run() 161{ 162 // Load images 163 left_src = imread(p.left); 164 right_src = imread(p.right); 165 if (left_src.empty()) throw runtime_error("can't open file \"" + p.left + "\""); 166 if (right_src.empty()) throw runtime_error("can't open file \"" + p.right + "\""); 167 cvtColor(left_src, left, COLOR_BGR2GRAY); 168 cvtColor(right_src, right, COLOR_BGR2GRAY); 169 d_left.upload(left); 170 d_right.upload(right); 171 172 imshow("left", left); 173 imshow("right", right); 174 175 // Set common parameters 176 bm = cuda::createStereoBM(p.ndisp); 177 bp = cuda::createStereoBeliefPropagation(p.ndisp); 178 csbp = cv::cuda::createStereoConstantSpaceBP(p.ndisp); 179 180 // Prepare disparity map of specified type 181 Mat disp(left.size(), CV_8U); 182 cuda::GpuMat d_disp(left.size(), CV_8U); 183 184 cout << endl; 185 printParams(); 186 187 running = true; 188 while (running) 189 { 190 workBegin(); 191 switch (p.method) 192 { 193 case Params::BM: 194 if (d_left.channels() > 1 || d_right.channels() > 1) 195 { 196 cout << "BM doesn't support color images\n"; 197 cvtColor(left_src, left, COLOR_BGR2GRAY); 198 cvtColor(right_src, right, COLOR_BGR2GRAY); 199 cout << "image_channels: " << left.channels() << endl; 200 d_left.upload(left); 201 d_right.upload(right); 202 imshow("left", left); 203 imshow("right", right); 204 } 205 bm->compute(d_left, d_right, d_disp); 206 break; 207 case Params::BP: bp->compute(d_left, d_right, d_disp); break; 208 case Params::CSBP: csbp->compute(d_left, d_right, d_disp); break; 209 } 210 workEnd(); 211 212 // Show results 213 d_disp.download(disp); 214 putText(disp, text(), Point(5, 25), FONT_HERSHEY_SIMPLEX, 1.0, Scalar::all(255)); 215 imshow("disparity", disp); 216 217 handleKey((char)waitKey(3)); 218 } 219} 220 221 222void App::printParams() const 223{ 224 cout << "--- Parameters ---\n"; 225 cout << "image_size: (" << left.cols << ", " << left.rows << ")\n"; 226 cout << "image_channels: " << left.channels() << endl; 227 cout << "method: " << p.method_str() << endl 228 << "ndisp: " << p.ndisp << endl; 229 switch (p.method) 230 { 231 case Params::BM: 232 cout << "win_size: " << bm->getBlockSize() << endl; 233 cout << "prefilter_sobel: " << bm->getPreFilterType() << endl; 234 break; 235 case Params::BP: 236 cout << "iter_count: " << bp->getNumIters() << endl; 237 cout << "level_count: " << bp->getNumLevels() << endl; 238 break; 239 case Params::CSBP: 240 cout << "iter_count: " << csbp->getNumIters() << endl; 241 cout << "level_count: " << csbp->getNumLevels() << endl; 242 break; 243 } 244 cout << endl; 245} 246 247 248void App::handleKey(char key) 249{ 250 switch (key) 251 { 252 case 27: 253 running = false; 254 break; 255 case 'p': case 'P': 256 printParams(); 257 break; 258 case 'g': case 'G': 259 if (left.channels() == 1 && p.method != Params::BM) 260 { 261 left = left_src; 262 right = right_src; 263 } 264 else 265 { 266 cvtColor(left_src, left, COLOR_BGR2GRAY); 267 cvtColor(right_src, right, COLOR_BGR2GRAY); 268 } 269 d_left.upload(left); 270 d_right.upload(right); 271 cout << "image_channels: " << left.channels() << endl; 272 imshow("left", left); 273 imshow("right", right); 274 break; 275 case 'm': case 'M': 276 switch (p.method) 277 { 278 case Params::BM: 279 p.method = Params::BP; 280 break; 281 case Params::BP: 282 p.method = Params::CSBP; 283 break; 284 case Params::CSBP: 285 p.method = Params::BM; 286 break; 287 } 288 cout << "method: " << p.method_str() << endl; 289 break; 290 case 's': case 'S': 291 if (p.method == Params::BM) 292 { 293 switch (bm->getPreFilterType()) 294 { 295 case 0: 296 bm->setPreFilterType(cv::StereoBM::PREFILTER_XSOBEL); 297 break; 298 case cv::StereoBM::PREFILTER_XSOBEL: 299 bm->setPreFilterType(0); 300 break; 301 } 302 cout << "prefilter_sobel: " << bm->getPreFilterType() << endl; 303 } 304 break; 305 case '1': 306 p.ndisp = p.ndisp == 1 ? 8 : p.ndisp + 8; 307 cout << "ndisp: " << p.ndisp << endl; 308 bm->setNumDisparities(p.ndisp); 309 bp->setNumDisparities(p.ndisp); 310 csbp->setNumDisparities(p.ndisp); 311 break; 312 case 'q': case 'Q': 313 p.ndisp = max(p.ndisp - 8, 1); 314 cout << "ndisp: " << p.ndisp << endl; 315 bm->setNumDisparities(p.ndisp); 316 bp->setNumDisparities(p.ndisp); 317 csbp->setNumDisparities(p.ndisp); 318 break; 319 case '2': 320 if (p.method == Params::BM) 321 { 322 bm->setBlockSize(min(bm->getBlockSize() + 1, 51)); 323 cout << "win_size: " << bm->getBlockSize() << endl; 324 } 325 break; 326 case 'w': case 'W': 327 if (p.method == Params::BM) 328 { 329 bm->setBlockSize(max(bm->getBlockSize() - 1, 2)); 330 cout << "win_size: " << bm->getBlockSize() << endl; 331 } 332 break; 333 case '3': 334 if (p.method == Params::BP) 335 { 336 bp->setNumIters(bp->getNumIters() + 1); 337 cout << "iter_count: " << bp->getNumIters() << endl; 338 } 339 else if (p.method == Params::CSBP) 340 { 341 csbp->setNumIters(csbp->getNumIters() + 1); 342 cout << "iter_count: " << csbp->getNumIters() << endl; 343 } 344 break; 345 case 'e': case 'E': 346 if (p.method == Params::BP) 347 { 348 bp->setNumIters(max(bp->getNumIters() - 1, 1)); 349 cout << "iter_count: " << bp->getNumIters() << endl; 350 } 351 else if (p.method == Params::CSBP) 352 { 353 csbp->setNumIters(max(csbp->getNumIters() - 1, 1)); 354 cout << "iter_count: " << csbp->getNumIters() << endl; 355 } 356 break; 357 case '4': 358 if (p.method == Params::BP) 359 { 360 bp->setNumLevels(bp->getNumLevels() + 1); 361 cout << "level_count: " << bp->getNumLevels() << endl; 362 } 363 else if (p.method == Params::CSBP) 364 { 365 csbp->setNumLevels(csbp->getNumLevels() + 1); 366 cout << "level_count: " << csbp->getNumLevels() << endl; 367 } 368 break; 369 case 'r': case 'R': 370 if (p.method == Params::BP) 371 { 372 bp->setNumLevels(max(bp->getNumLevels() - 1, 1)); 373 cout << "level_count: " << bp->getNumLevels() << endl; 374 } 375 else if (p.method == Params::CSBP) 376 { 377 csbp->setNumLevels(max(csbp->getNumLevels() - 1, 1)); 378 cout << "level_count: " << csbp->getNumLevels() << endl; 379 } 380 break; 381 } 382} 383