samples/gpu/brox_optical_flow.cpp

   1 #include <iostream>
   2 #include <iomanip>
   3 #include <string>
   4
   5 #include "cvconfig.h"
   6 #include "opencv2/core/core.hpp"
   7 #include "opencv2/highgui/highgui.hpp"
   8 #include "opencv2/gpu/gpu.hpp"
   9
  10 using namespace std;
  11 using namespace cv;
  12 using namespace cv::gpu;
  13
  14 void getFlowField(const Mat& u, const Mat& v, Mat& flowField);
  15
  16 int main(int argc, const char* argv[])
  17 {
  18     try
  19     {
  20         const char* keys =
  21            "{ h  | help      | false | print help message }"
  22            "{ l  | left      |       | specify left image }"
  23            "{ r  | right     |       | specify right image }"
  24            "{ s  | scale     | 0.8   | set pyramid scale factor }"
  25            "{ a  | alpha     | 0.197 | set alpha }"
  26            "{ g  | gamma     | 50.0  | set gamma }"
  27            "{ i  | inner     | 10    | set number of inner iterations }"
  28            "{ o  | outer     | 77    | set number of outer iterations }"
  29            "{ si | solver    | 10    | set number of basic solver iterations }"
  30            "{ t  | time_step | 0.1   | set frame interpolation time step }";
  31
  32         CommandLineParser cmd(argc, argv, keys);
  33
  34         if (cmd.get<bool>("help"))
  35         {
  36             cout << "Usage: brox_optical_flow [options]" << endl;
  37             cout << "Avaible options:" << endl;
  38             cmd.printParams();
  39             return 0;
  40         }
  41
  42         string frame0Name = cmd.get<string>("left");
  43         string frame1Name = cmd.get<string>("right");
  44         float scale = cmd.get<float>("scale");
  45         float alpha = cmd.get<float>("alpha");
  46         float gamma = cmd.get<float>("gamma");
  47         int inner_iterations = cmd.get<int>("inner");
  48         int outer_iterations = cmd.get<int>("outer");
  49         int solver_iterations = cmd.get<int>("solver");
  50         float timeStep = cmd.get<float>("time_step");
  51
  52         if (frame0Name.empty() || frame1Name.empty())
  53         {
  54             cerr << "Missing input file names" << endl;
  55             return -1;
  56         }
  57
  58         Mat frame0Color = imread(frame0Name);
  59         Mat frame1Color = imread(frame1Name);
  60
  61         if (frame0Color.empty() || frame1Color.empty())
  62         {
  63             cout << "Can't load input images" << endl;
  64             return -1;
  65         }
  66
  67         cv::gpu::printShortCudaDeviceInfo(cv::gpu::getDevice());
  68
  69         cout << "OpenCV / NVIDIA Computer Vision" << endl;
  70         cout << "Optical Flow Demo: Frame Interpolation" << endl;
  71         cout << "=========================================" << endl;
  72
  73         namedWindow("Forward flow");
  74         namedWindow("Backward flow");
  75
  76         namedWindow("Interpolated frame");
  77
  78         cout << "Press:" << endl;
  79         cout << "\tESC to quit" << endl;
  80         cout << "\t'a' to move to the previous frame" << endl;
  81         cout << "\t's' to move to the next frame\n" << endl;
  82
  83         frame0Color.convertTo(frame0Color, CV_32F, 1.0 / 255.0);
  84         frame1Color.convertTo(frame1Color, CV_32F, 1.0 / 255.0);
  85
  86         Mat frame0Gray, frame1Gray;
  87
  88         cvtColor(frame0Color, frame0Gray, COLOR_BGR2GRAY);
  89         cvtColor(frame1Color, frame1Gray, COLOR_BGR2GRAY);
  90
  91         GpuMat d_frame0(frame0Gray);
  92         GpuMat d_frame1(frame1Gray);
  93
  94         cout << "Estimating optical flow" << endl;
  95
  96         BroxOpticalFlow d_flow(alpha, gamma, scale, inner_iterations, outer_iterations, solver_iterations);
  97
  98         cout << "\tForward..." << endl;
  99
 100         GpuMat d_fu, d_fv;
 101
 102         d_flow(d_frame0, d_frame1, d_fu, d_fv);
 103
 104         Mat flowFieldForward;
 105         getFlowField(Mat(d_fu), Mat(d_fv), flowFieldForward);
 106
 107         cout << "\tBackward..." << endl;
 108
 109         GpuMat d_bu, d_bv;
 110
 111         d_flow(d_frame1, d_frame0, d_bu, d_bv);
 112
 113         Mat flowFieldBackward;
 114         getFlowField(Mat(d_bu), Mat(d_bv), flowFieldBackward);
 115
 116         cout << "Interpolating..." << endl;
 117
 118         // first frame color components
 119         GpuMat d_b, d_g, d_r;
 120
 121         // second frame color components
 122         GpuMat d_bt, d_gt, d_rt;
 123
 124         // prepare color components on host and copy them to device memory
 125         Mat channels[3];
 126         cv::split(frame0Color, channels);
 127
 128         d_b.upload(channels[0]);
 129         d_g.upload(channels[1]);
 130         d_r.upload(channels[2]);
 131
 132         cv::split(frame1Color, channels);
 133
 134         d_bt.upload(channels[0]);
 135         d_gt.upload(channels[1]);
 136         d_rt.upload(channels[2]);
 137
 138         // temporary buffer
 139         GpuMat d_buf;
 140
 141         // intermediate frame color components (GPU memory)
 142         GpuMat d_rNew, d_gNew, d_bNew;
 143
 144         GpuMat d_newFrame;
 145
 146         vector<Mat> frames;
 147         frames.reserve(static_cast<int>(1.0f / timeStep) + 2);
 148
 149         frames.push_back(frame0Color);
 150
 151         // compute interpolated frames
 152         for (float timePos = timeStep; timePos < 1.0f; timePos += timeStep)
 153         {
 154             // interpolate blue channel
 155             interpolateFrames(d_b, d_bt, d_fu, d_fv, d_bu, d_bv, timePos, d_bNew, d_buf);
 156
 157             // interpolate green channel
 158             interpolateFrames(d_g, d_gt, d_fu, d_fv, d_bu, d_bv, timePos, d_gNew, d_buf);
 159
 160             // interpolate red channel
 161             interpolateFrames(d_r, d_rt, d_fu, d_fv, d_bu, d_bv, timePos, d_rNew, d_buf);
 162
 163             GpuMat channels3[] = {d_bNew, d_gNew, d_rNew};
 164             merge(channels3, 3, d_newFrame);
 165
 166             frames.push_back(Mat(d_newFrame));
 167
 168             cout << setprecision(4) << timePos * 100.0f << "%\r";
 169         }
 170
 171         frames.push_back(frame1Color);
 172
 173         cout << setw(5) << "100%" << endl;
 174
 175         cout << "Done" << endl;
 176
 177         imshow("Forward flow", flowFieldForward);
 178         imshow("Backward flow", flowFieldBackward);
 179
 180         int currentFrame = 0;
 181
 182         imshow("Interpolated frame", frames[currentFrame]);
 183
 184         for(;;)
 185         {
 186             int key = toupper(waitKey(10) & 0xff);
 187
 188             switch (key)
 189             {
 190             case 27:
 191                 return 0;
 192
 193             case 'A':
 194                 if (currentFrame > 0)
 195                     --currentFrame;
 196
 197                 imshow("Interpolated frame", frames[currentFrame]);
 198                 break;
 199
 200             case 'S':
 201                 if (currentFrame < static_cast<int>(frames.size()) - 1)
 202                     ++currentFrame;
 203
 204                 imshow("Interpolated frame", frames[currentFrame]);
 205                 break;
 206             }
 207         }
 208     }
 209     catch (const exception& ex)
 210     {
 211         cerr << ex.what() << endl;
 212         return -1;
 213     }
 214     catch (...)
 215     {
 216         cerr << "Unknow error" << endl;
 217         return -1;
 218     }
 219 }
 220
 221 template <typename T> inline T clamp (T x, T a, T b)
 222 {
 223     return ((x) > (a) ? ((x) < (b) ? (x) : (b)) : (a));
 224 }
 225
 226 template <typename T> inline T mapValue(T x, T a, T b, T c, T d)
 227 {
 228     x = clamp(x, a, b);
 229     return c + (d - c) * (x - a) / (b - a);
 230 }
 231
 232 void getFlowField(const Mat& u, const Mat& v, Mat& flowField)
 233 {
 234     float maxDisplacement = 1.0f;
 235
 236     for (int i = 0; i < u.rows; ++i)
 237     {
 238         const float* ptr_u = u.ptr<float>(i);
 239         const float* ptr_v = v.ptr<float>(i);
 240
 241         for (int j = 0; j < u.cols; ++j)
 242         {
 243             float d = max(fabsf(ptr_u[j]), fabsf(ptr_v[j]));
 244
 245             if (d > maxDisplacement)
 246                 maxDisplacement = d;
 247         }
 248     }
 249
 250     flowField.create(u.size(), CV_8UC4);
 251
 252     for (int i = 0; i < flowField.rows; ++i)
 253     {
 254         const float* ptr_u = u.ptr<float>(i);
 255         const float* ptr_v = v.ptr<float>(i);
 256
 257
 258         Vec4b* row = flowField.ptr<Vec4b>(i);
 259
 260         for (int j = 0; j < flowField.cols; ++j)
 261         {
 262             row[j][0] = 0;
 263             row[j][1] = static_cast<unsigned char> (mapValue (-ptr_v[j], -maxDisplacement, maxDisplacement, 0.0f, 255.0f));
 264             row[j][2] = static_cast<unsigned char> (mapValue ( ptr_u[j], -maxDisplacement, maxDisplacement, 0.0f, 255.0f));
 265             row[j][3] = 255;
 266         }
 267     }
 268 }