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