Optical flow video tracking with per-triplet consistency check, using calcOpticalFlowFarneback from OpenCV

CMakeLists.txt

cmake_minimum_required(VERSION 3.13)
project(video_optflow)

set(CMAKE_CXX_STANDARD 14)

find_package(OpenMP)
if (OpenMP_CXX_FOUND)
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
    set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}")
else()
    message(WARNING "OpenMP not found!")
endif()

find_package(OpenCV 3.3 REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})

add_executable(video_optflow main.cpp)
target_link_libraries(video_optflow ${OpenCV_LIBS})

main.cpp

#include <iostream>

#ifdef _OPENMP
#include <omp.h>
#endif

#include "opencv2/imgproc.hpp"
#include "opencv2/videoio.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/video/tracking.hpp"


template <typename T>
std::string to_string(T value)
{
	std::ostringstream ss;
	ss << value;
	return ss.str();
}


void check(bool statement, const std::string &message, int line)
{
	if (!statement) {
		throw std::runtime_error(message + " (at line " + to_string(line) + ")");
	}
}

#define CHECK(statement, message) check(statement, message, __LINE__)


void optflow_with_lr_check(const cv::Mat &grayframe0, const cv::Mat &grayframe1, const cv::Mat &grayframe2,
		std::vector<cv::Vec2i> &points0, std::vector<cv::Vec2i> &points1, std::vector<cv::Vec2i> &points2,
		float &resAvgFlow, float &resExtremumFlow,
		bool verbose=false)
{
	const float threshold = 0.25f;

	cv::Mat flow01, flow12, flow20;

	#pragma omp parallel sections
	{
		#pragma omp section
		{
			cv::calcOpticalFlowFarneback(grayframe0, grayframe1, flow01,
										 0.5, 5, 6, 3,
										 6, 1.2, 0);
		}
		#pragma omp section
		{
			cv::calcOpticalFlowFarneback(grayframe1, grayframe2, flow12,
										 0.5, 5, 6, 3,
										 6, 1.2, 0);
		}
		#pragma omp section
		{
			cv::calcOpticalFlowFarneback(grayframe2, grayframe0, flow20,
										 0.5, 5, 6, 3,
										 6, 1.2, 0);
		}
	}

	std::vector<float> norms;

	float normThreshold = -1.0f;
	float normExtremum = -1.0f;
	float avgFlow = 0.0f;
	for (int stage = 0; stage < 2; ++stage) {
		#pragma omp parallel for
		for (ptrdiff_t j0 = 0; j0 < flow01.rows; ++j0) {
			const float* row01 = flow01.ptr<float>(j0);
			for (ptrdiff_t i0 = 0; i0 < flow01.cols; ++i0) {
				float dx01 = row01[2 * i0 + 0];
				float dy01 = row01[2 * i0 + 1];
				if (fabs(dy01) > std::max(grayframe0.rows, grayframe0.cols)) continue;
				if (fabs(dx01) > std::max(grayframe0.rows, grayframe0.cols)) continue;
				ptrdiff_t j1 = std::round(j0 + 0.5f + dy01);
				ptrdiff_t i1 = std::round(i0 + 0.5f + dx01);

				if (j1 < 0 || j1 >= grayframe0.rows) continue;
				if (i1 < 0 || i1 >= grayframe0.cols) continue;

				const float* row12 = flow12.ptr<float>(j1);
				float dx12 = row12[2 * i1 + 0];
				float dy12 = row12[2 * i1 + 1];
				ptrdiff_t j2 = std::round(j0 + 0.5f + dy01 + dy12);
				ptrdiff_t i2 = std::round(i0 + 0.5f + dx01 + dx12);

				if (j2 < 0 || j2 >= grayframe0.rows) continue;
				if (i2 < 0 || i2 >= grayframe0.cols) continue;

				const float* row20 = flow20.ptr<float>(j2);
				float dx20 = row20[2 * i2 + 0];
				float dy20 = row20[2 * i2 + 1];

				ptrdiff_t j20 = std::round(j0 + 0.5f + dy01 + dy12 + dy20);
				ptrdiff_t i20 = std::round(i0 + 0.5f + dx01 + dx12 + dx20);

				if (j20 < 0 || j20 >= grayframe0.rows) continue;
				if (i20 < 0 || i20 >= grayframe0.cols) continue;

				float dx02 = dx01 + dx12;
				float dy02 = dy01 + dy12;

				float dxMotionDelta = fabs(dx01 - dx12);
				float dyMotionDelta = fabs(dy01 - dy12);

				float dxError = fabs(dx02 + dx20);
				float dyError = fabs(dy02 + dy20);

				if (abs(j20 - j0) <= 1 && abs(i20 - i0) <= 1 &&
//					dxMotionDelta * dxMotionDelta + dyMotionDelta * dyMotionDelta < 0.5f*std::max(dx01*dx01+dy01*dy01, dx12*dx12+dy12*dy12) &&
					dxError * dxError + dyError * dyError < threshold * threshold) {
					float norm = sqrtf(dx20 * dx20 + dy20 * dy20);
					if (stage == 0) {
						#pragma omp critical
						{
							norms.push_back(norm);
						}
					} else {
						CHECK(normThreshold >= 0.0f, "Norm threshold not initialized!");
						if (norm > normThreshold) {
							#pragma omp critical
							{
								points0.push_back(cv::Vec2i(i0, j0));
								points1.push_back(cv::Vec2i(i1, j1));
								points2.push_back(cv::Vec2i(i2, j2));
								avgFlow += sqrtf(norm);
							}
						}
					}
				}
			}
		}
		if (stage == 0) {
			std::sort(norms.begin(), norms.end());
			int tooSmall = 0;
			int tooBig = 0;
			const int max = 45;
			std::vector<int> votes(max, 0);
			for (int i = 0; i < norms.size(); ++i) {
				float norm = norms[i];
				int bin = std::floor(norm);
				if (bin < 0) {
					++tooSmall;
				} else if (bin >= max) {
					++tooBig;
				} else {
					++votes[bin];
				}
			}
			if (verbose) {
				std::cout << "Vectors length bins: ";
				for (int i = 0; i < max; ++i) {
					std::cout << votes[i];
					if (i + 1 < max) std::cout << ", ";
				}
				std::cout << " (and: " << tooSmall << " too small/" << tooBig << " too big)" << std::endl;
			}

			std::vector<float> smoothedVotes(max, 0.0f);
			for (int i = 0; i < max; ++i) {
				float v = votes[i];
				float w = 1.0f;
				if (i + 1 < max) { v += 0.5f * votes[i + 1]; w += 0.5f;}
				if (i - 1 >= 0)  { v += 0.5f * votes[i - 1]; w += 0.5f;}
				smoothedVotes[i] = v / w;
			}

			int extremum = -1;
			float extremumVotes = -1.0f;
			for (int i = max - 1; i >= 0; --i) {
				bool localMax = true;
				if (smoothedVotes[i] < 100) { localMax = false; }
				if (i + 1 < max) { if (smoothedVotes[i + 1] > smoothedVotes[i]) localMax = false; }
				if (i - 1 >= 0)  { if (smoothedVotes[i - 1] > smoothedVotes[i]) localMax = false; }
				if (localMax) {
					extremum = i;
					extremumVotes = smoothedVotes[extremum];
					break;
				}
			}
//			CHECK(extremum != -1, "Extremum not found!");
//			CHECK(extremumVotes > tooBig, "Range of bins is too small?!");

			normExtremum = extremum;
			normThreshold = 0.5f;
		}
	}
	resAvgFlow = avgFlow / points0.size();
	resExtremumFlow = normExtremum;
}

cv::Mat drawPoints(const std::vector<unsigned int> &indices, const int limitPoints,
		const cv::Mat &image, std::vector<cv::Vec2i> points0, std::vector<cv::Vec2i> points1=std::vector<cv::Vec2i>())
{
	const cv::Scalar blue = cv::Scalar(255, 0, 0);
	const cv::Scalar red  = cv::Scalar(0, 0, 255);
	const int radius = 4;

	cv::Mat res;
	image.copyTo(res);

	if (limitPoints != 0 && limitPoints < points0.size() && indices.size() > 0) {
		{
			std::vector<cv::Vec2i> tmp0 = points0;
			points0.resize(limitPoints);
			for (int i = 0; i < limitPoints; ++i) {
				points0[i] = tmp0[indices[i]];
			}
		}
		if (points1.size() > 0) {
			std::vector<cv::Vec2i> tmp1 = points1;
			points1.resize(limitPoints);
			for (int i = 0; i < limitPoints; ++i) {
				points1[i] = tmp1[indices[i]];
			}
		}
	}

	if (points1.size() > 0) {
		for (ptrdiff_t i = 0; i < points0.size(); ++i) {
			int x0 = points0[i][0];
			int y0 = points0[i][1];
			int x1 = points1[i][0];
			int y1 = points1[i][1];
			cv::line(res, cv::Point(x0, y0), cv::Point(x1, y1), blue);
		}
	}

	for (ptrdiff_t i = 0; i < points0.size(); ++i) {
		int x0 = points0[i][0];
		int y0 = points0[i][1];
		cv::circle(res, cv::Point(x0, y0), radius, red);
	}

	return res;
}

cv::Mat toGray(const cv::Mat &bgr)
{
	cv::Mat gray;
	cv::cvtColor(bgr, gray, cv::COLOR_BGR2GRAY);
	return gray;
}


int main()
{
	const int KEY_ESCAPE = 27;
	const int KEY_SPACE  = 32;

	const int minStep = 1;
	const int maxStep = 30;
	const float avgFlowMin = 5.0f;
	const float avgFlowMax = 3 * avgFlowMin;
	const float extremumFlowMin = avgFlowMin;
	const float extremumFlowMax = 1.5f * avgFlowMax;
	const int downscale = 4;

	const int limitPointsPreview = 2391;

	// youtube-dl -f bestvideo https://www.youtube.com/watch?v=GTFICUsvqi0
	const std::string videoFile = "Dover castle DJI Phantom 3 pro drone-GTFICUsvqi0.webm";

	cv::VideoCapture cap;
	cap.open(videoFile);
	CHECK(cap.isOpened(), "Can't open video file!");

	int step = (minStep + maxStep) / 2;
	bool resolutionPrinted = false;

	cv::Mat frame0;
	cv::Mat frame1;
	cv::Mat frame2;
	cv::Mat nextFrame;
	int frameIndex = 0;
	while (true) {
		for (int i = 0; i < step; ++i) {
			cap >> nextFrame;
			++frameIndex;
		}
		if (nextFrame.empty()) break;

		if (!resolutionPrinted) {
			std::cout << "Resolution: " << nextFrame.cols << "x" << nextFrame.rows;
		}
		if (downscale > 1) {
			int cur = 1;
			while (cur < downscale) {
				cv::Mat downscaled;
				cv::pyrDown(nextFrame, downscaled);
				nextFrame = downscaled;
				cur *= 2;
			}
			if (!resolutionPrinted) {
				std::cout << " (downscaled to " << nextFrame.cols << "x" << nextFrame.rows << ")";
			}
		}
		if (!resolutionPrinted) {
			std::cout << std::endl;
			resolutionPrinted = true;
		}

		std::swap(frame0, frame1);
		std::swap(frame1, frame2);
		nextFrame.copyTo(frame2);
		if (frame0.empty()) continue;

		std::vector<cv::Vec2i> points0;
		std::vector<cv::Vec2i> points1;
		std::vector<cv::Vec2i> points2;

		float avgFlow;
		float extremumFlow;
		optflow_with_lr_check(toGray(frame0), toGray(frame1), toGray(frame2), points0, points1, points2,
				avgFlow, extremumFlow, false);

		std::cout << "Frame #" << frameIndex << ": " << points0.size() << " points (avgFlow=" << avgFlow << " extremumFlow=" << extremumFlow << ")" << std::endl;
		std::vector<unsigned int> indices(points0.size());
		for (int i = 0; i < points0.size(); ++i) {
			indices[i] = i;
		}
		std::random_shuffle(indices.begin(), indices.end());
		cv::imshow("Frame 0", drawPoints(indices, limitPointsPreview, frame0, points0));
		cv::imshow("Frame 1", drawPoints(indices, limitPointsPreview, frame1, points1));
		cv::imshow("Frame 2", drawPoints(indices, limitPointsPreview, frame2, points2, points0));

		if ((avgFlow > avgFlowMax || points0.size() < 20000) && step > minStep) {
			step = minStep;
			std::cout << "Slow down to " << step << "!" << std::endl;
		} else if (avgFlow < avgFlowMin && step < maxStep) {
			step = std::min(step + 4, maxStep);
			std::cout << "Speed up to " << step << "!" << std::endl;
		}

		int key = cv::waitKey(10);
		if (key == KEY_SPACE) {
			std::cout << "Paused!" << std::endl;
			while (cv::waitKey(10) != KEY_SPACE) {}
		} else if (key == KEY_ESCAPE) {
			break;
		}
	}

	return 0;
}