TVL1 optical flow example moved from opencv to opencv_contrib

Author: vpisarev

modules/optflow/samples/tvl1_optical_flow.cpp

@@ -0,0 +1,206 @@
+#include <iostream>
+#include <fstream>
+
+#include <opencv2/core/utility.hpp>
+#include "opencv2/video.hpp"
+#include "opencv2/optflow.hpp"
+#include "opencv2/imgcodecs.hpp"
+#include "opencv2/highgui.hpp"
+
+using namespace cv;
+using namespace std;
+using namespace optflow;
+
+inline bool isFlowCorrect(Point2f u)
+{
+    return !cvIsNaN(u.x) && !cvIsNaN(u.y) && fabs(u.x) < 1e9 && fabs(u.y) < 1e9;
+}
+
+static Vec3b computeColor(float fx, float fy)
+{
+    static bool first = true;
+
+    // relative lengths of color transitions:
+    // these are chosen based on perceptual similarity
+    // (e.g. one can distinguish more shades between red and yellow
+    //  than between yellow and green)
+    const int RY = 15;
+    const int YG = 6;
+    const int GC = 4;
+    const int CB = 11;
+    const int BM = 13;
+    const int MR = 6;
+    const int NCOLS = RY + YG + GC + CB + BM + MR;
+    static Vec3i colorWheel[NCOLS];
+
+    if (first)
+    {
+        int k = 0;
+
+        for (int i = 0; i < RY; ++i, ++k)
+            colorWheel[k] = Vec3i(255, 255 * i / RY, 0);
+
+        for (int i = 0; i < YG; ++i, ++k)
+            colorWheel[k] = Vec3i(255 - 255 * i / YG, 255, 0);
+
+        for (int i = 0; i < GC; ++i, ++k)
+            colorWheel[k] = Vec3i(0, 255, 255 * i / GC);
+
+        for (int i = 0; i < CB; ++i, ++k)
+            colorWheel[k] = Vec3i(0, 255 - 255 * i / CB, 255);
+
+        for (int i = 0; i < BM; ++i, ++k)
+            colorWheel[k] = Vec3i(255 * i / BM, 0, 255);
+
+        for (int i = 0; i < MR; ++i, ++k)
+            colorWheel[k] = Vec3i(255, 0, 255 - 255 * i / MR);
+
+        first = false;
+    }
+
+    const float rad = sqrt(fx * fx + fy * fy);
+    const float a = atan2(-fy, -fx) / (float)CV_PI;
+
+    const float fk = (a + 1.0f) / 2.0f * (NCOLS - 1);
+    const int k0 = static_cast<int>(fk);
+    const int k1 = (k0 + 1) % NCOLS;
+    const float f = fk - k0;
+
+    Vec3b pix;
+
+    for (int b = 0; b < 3; b++)
+    {
+        const float col0 = colorWheel[k0][b] / 255.f;
+        const float col1 = colorWheel[k1][b] / 255.f;
+
+        float col = (1 - f) * col0 + f * col1;
+
+        if (rad <= 1)
+            col = 1 - rad * (1 - col); // increase saturation with radius
+        else
+            col *= .75; // out of range
+
+        pix[2 - b] = static_cast<uchar>(255.f * col);
+    }
+
+    return pix;
+}
+
+static void drawOpticalFlow(const Mat_<Point2f>& flow, Mat& dst, float maxmotion = -1)
+{
+    dst.create(flow.size(), CV_8UC3);
+    dst.setTo(Scalar::all(0));
+
+    // determine motion range:
+    float maxrad = maxmotion;
+
+    if (maxmotion <= 0)
+    {
+        maxrad = 1;
+        for (int y = 0; y < flow.rows; ++y)
+        {
+            for (int x = 0; x < flow.cols; ++x)
+            {
+                Point2f u = flow(y, x);
+
+                if (!isFlowCorrect(u))
+                    continue;
+
+                maxrad = max(maxrad, sqrt(u.x * u.x + u.y * u.y));
+            }
+        }
+    }
+
+    for (int y = 0; y < flow.rows; ++y)
+    {
+        for (int x = 0; x < flow.cols; ++x)
+        {
+            Point2f u = flow(y, x);
+
+            if (isFlowCorrect(u))
+                dst.at<Vec3b>(y, x) = computeColor(u.x / maxrad, u.y / maxrad);
+        }
+    }
+}
+
+// binary file format for flow data specified here:
+// http://vision.middlebury.edu/flow/data/
+static void writeOpticalFlowToFile(const Mat_<Point2f>& flow, const string& fileName)
+{
+    static const char FLO_TAG_STRING[] = "PIEH";
+
+    ofstream file(fileName.c_str(), ios_base::binary);
+
+    file << FLO_TAG_STRING;
+
+    file.write((const char*) &flow.cols, sizeof(int));
+    file.write((const char*) &flow.rows, sizeof(int));
+
+    for (int i = 0; i < flow.rows; ++i)
+    {
+        for (int j = 0; j < flow.cols; ++j)
+        {
+            const Point2f u = flow(i, j);
+
+            file.write((const char*) &u.x, sizeof(float));
+            file.write((const char*) &u.y, sizeof(float));
+        }
+    }
+}
+
+int main(int argc, const char* argv[])
+{
+    cv::CommandLineParser parser(argc, argv, "{help h || show help message}"
+            "{ @frame0 | | frame 0}{ @frame1 | | frame 1}{ @output | | output flow}");
+    if (parser.has("help"))
+    {
+        parser.printMessage();
+        return 0;
+    }
+    string frame0_name = parser.get<string>("@frame0");
+    string frame1_name = parser.get<string>("@frame1");
+    string file = parser.get<string>("@output");
+    if (frame0_name.empty() || frame1_name.empty() || file.empty())
+    {
+        cerr << "Usage : " << argv[0] << " [<frame0>] [<frame1>] [<output_flow>]" << endl;
+        return -1;
+    }
+
+    Mat frame0 = imread(frame0_name, IMREAD_GRAYSCALE);
+    Mat frame1 = imread(frame1_name, IMREAD_GRAYSCALE);
+
+    if (frame0.empty())
+    {
+        cerr << "Can't open image ["  << parser.get<string>("frame0") << "]" << endl;
+        return -1;
+    }
+    if (frame1.empty())
+    {
+        cerr << "Can't open image ["  << parser.get<string>("frame1") << "]" << endl;
+        return -1;
+    }
+
+    if (frame1.size() != frame0.size())
+    {
+        cerr << "Images should be of equal sizes" << endl;
+        return -1;
+    }
+
+    Mat_<Point2f> flow;
+    Ptr<DualTVL1OpticalFlow> tvl1 = DualTVL1OpticalFlow::create();
+
+    const double start = (double)getTickCount();
+    tvl1->calc(frame0, frame1, flow);
+    const double timeSec = (getTickCount() - start) / getTickFrequency();
+    cout << "calcOpticalFlowDual_TVL1 : " << timeSec << " sec" << endl;
+
+    Mat out;
+    drawOpticalFlow(flow, out);
+    if (!file.empty())
+        writeOpticalFlowToFile(flow, file);
+
+    imshow("Flow", out);
+    waitKey();
+
+    return 0;
+}