File: optical_flow.cc

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (85 lines) | stat: -rw-r--r-- 2,448 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#include <caffe2/video/optical_flow.h>

namespace caffe2 {

void OpticalFlowExtractor(
    const cv::Mat& prev_gray,
    const cv::Mat& curr_gray,
    const int flow_alg_type,
    cv::Mat& flow) {
#if CV_MAJOR_VERSION >= 4
  cv::Ptr<cv::DISOpticalFlow> tvl1 = cv::DISOpticalFlow::create();
#else
  cv::Ptr<cv::DualTVL1OpticalFlow> tvl1 = cv::DualTVL1OpticalFlow::create();
#endif
  switch (flow_alg_type) {
    case FLowAlgType::FarnebackOpticalFlow:
      cv::calcOpticalFlowFarneback(
          prev_gray,
          curr_gray,
          flow,
          std::sqrt(2) / 2.0,
          5,
          10,
          2,
          7,
          1.5,
          cv::OPTFLOW_FARNEBACK_GAUSSIAN);
      break;
    case FLowAlgType::DensePyrLKOpticalFlow:
      LOG(ERROR) << "DensePyrLKOpticalFlow only has sparse version on CPU";
      break;
    case FLowAlgType::BroxOpticalFlow:
      LOG(ERROR) << "BroxOpticalFlow on CPU is not available";
      break;
    case FLowAlgType::OpticalFlowDual_TVL1:
      tvl1->calc(prev_gray, curr_gray, flow);
      break;
    default:
      LOG(ERROR) << "Unsupported optical flow type " << flow_alg_type;
      break;
  }
}

void MergeOpticalFlow(cv::Mat& prev_flow, const cv::Mat& curr_flow) {
  const int rows = prev_flow.rows;
  const int cols = prev_flow.cols;

  // merge two optical flows into one
  for (int y = 0; y < rows; y++) {
    for (int x = 0; x < cols; x++) {
      cv::Point2f u = prev_flow.at<cv::Point2f>(y, x);
      // get the new location
      int x_new = std::min(cols - 1, std::max(0, cvRound(u.x + x)));
      int y_new = std::min(rows - 1, std::max(0, cvRound(u.y + y)));
      cv::Point2f u_new = curr_flow.at<cv::Point2f>(y_new, x_new);

      // update the flow
      prev_flow.at<cv::Point2f>(y, x) += u_new;
    }
  }
}

void MultiFrameOpticalFlowExtractor(
    const std::vector<cv::Mat>& grays,
    const int optical_flow_alg_type,
    cv::Mat& flow) {
  int num_frames = grays.size();
  CAFFE_ENFORCE_GE(num_frames, 2, "need at least 2 frames!");

  // compute optical flow for every two frames
  std::vector<cv::Mat> flows;
  for (int i = 0; i < num_frames - 1; i++) {
    cv::Mat tmp;
    OpticalFlowExtractor(grays[i], grays[i + 1], optical_flow_alg_type, tmp);
    flows.push_back(tmp);
  }

  flows[0].copyTo(flow);
  // aggregate optical flow across multiple frame
  for (int i = 1; i < num_frames - 1; i++) {
    MergeOpticalFlow(flow, flows[i]);
  }
}

} // namespace caffe2