mirror/opencv
1
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2025-06-08 01:53:19 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #24913 from usyntest:optical-flow-sample-raft

Browse Source 
Raft support added in this sample code #24913

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake

fix: https://github.com/opencv/opencv/issues/24424 Update DNN Optical Flow sample with RAFT model
I implemented both RAFT and FlowNet v2 leaving it to the user which one he wants to use to estimate the optical flow.

Co-authored-by: Uday Sharma <uday@192.168.1.35>
This commit is contained in:
uday 2024-01-29 20:07:52 +05:30 committed by GitHub
parent 73acf08844
commit 03994163b5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 43 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
samples/dnn/optical_flow.py
View File

@ -1,13 +1,19 @@
#!/usr/bin/env python #!/usr/bin/env python
''' '''
This sample using FlowNet v2 model to calculate optical flow. This sample using FlowNet v2 and RAFT model to calculate optical flow.
Original paper: https://arxiv.org/abs/1612.01925.
Original repo: https://github.com/lmb-freiburg/flownet2. FlowNet v2 Original Paper: https://arxiv.org/abs/1612.01925.
FlowNet v2 Repo: https://github.com/lmb-freiburg/flownet2.
Download the converted .caffemodel model from https://drive.google.com/open?id=16qvE9VNmU39NttpZwZs81Ga8VYQJDaWZ Download the converted .caffemodel model from https://drive.google.com/open?id=16qvE9VNmU39NttpZwZs81Ga8VYQJDaWZ
and .prototxt from https://drive.google.com/file/d/1RyNIUsan1ZOh2hpYIH36A-jofAvJlT6a/view?usp=sharing. and .prototxt from https://drive.google.com/file/d/1RyNIUsan1ZOh2hpYIH36A-jofAvJlT6a/view?usp=sharing.
Otherwise download original model from https://lmb.informatik.uni-freiburg.de/resources/binaries/flownet2/flownet2-models.tar.gz, Otherwise download original model from https://lmb.informatik.uni-freiburg.de/resources/binaries/flownet2/flownet2-models.tar.gz,
convert .h5 model to .caffemodel and modify original .prototxt using .prototxt from link above. convert .h5 model to .caffemodel and modify original .prototxt using .prototxt from link above.
RAFT Original Paper: https://arxiv.org/pdf/2003.12039.pdf
RAFT Repo: https://github.com/princeton-vl/RAFT
Download the .onnx model from here https://github.com/opencv/opencv_zoo/raw/281d232cd99cd920853106d853c440edd35eb442/models/optical_flow_estimation_raft/optical_flow_estimation_raft_2023aug.onnx.
''' '''
import argparse import argparse
@ -17,8 +23,11 @@ import cv2 as cv
class OpticalFlow(object): class OpticalFlow(object):
def __init__(self, proto, model, height, width): def __init__(self, model, height, width, proto=""):
self.net = cv.dnn.readNetFromCaffe(proto, model) if proto:
self.net = cv.dnn.readNetFromCaffe(proto, model)
else:
self.net = cv.dnn.readNet(model)
self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV) self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
self.height = height self.height = height
self.width = width self.width = width
@ -26,8 +35,10 @@ class OpticalFlow(object):
def compute_flow(self, first_img, second_img): def compute_flow(self, first_img, second_img):
inp0 = cv.dnn.blobFromImage(first_img, size=(self.width, self.height)) inp0 = cv.dnn.blobFromImage(first_img, size=(self.width, self.height))
inp1 = cv.dnn.blobFromImage(second_img, size=(self.width, self.height)) inp1 = cv.dnn.blobFromImage(second_img, size=(self.width, self.height))
self.net.setInputsNames(["img0", "img1"])
self.net.setInput(inp0, "img0") self.net.setInput(inp0, "img0")
self.net.setInput(inp1, "img1") self.net.setInput(inp1, "img1")
flow = self.net.forward() flow = self.net.forward()
output = self.motion_to_color(flow) output = self.motion_to_color(flow)
return output return output
@ -46,7 +57,7 @@ class OpticalFlow(object):
rad = rad[..., np.newaxis] / maxrad rad = rad[..., np.newaxis] / maxrad
a = np.arctan2(-fy / maxrad, -fx / maxrad) / np.pi a = np.arctan2(-fy / maxrad, -fx / maxrad) / np.pi
fk = (a + 1) / 2.0 * (ncols - 1) fk = (a + 1) / 2.0 * (ncols - 1)
k0 = fk.astype(np.int) k0 = fk.astype(np.int32)
k1 = (k0 + 1) % ncols k1 = (k0 + 1) % ncols
f = fk[..., np.newaxis] - k0[..., np.newaxis] f = fk[..., np.newaxis] - k0[..., np.newaxis]
@ -59,41 +70,47 @@ class OpticalFlow(object):
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Use this script to calculate optical flow using FlowNetv2', parser = argparse.ArgumentParser(description='Use this script to calculate optical flow',
formatter_class=argparse.ArgumentDefaultsHelpFormatter) formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-input', '-i', required=True, help='Path to input video file. Skip this argument to capture frames from a camera.') parser.add_argument('-input', '-i', required=True, help='Path to input video file. Skip this argument to capture frames from a camera.')
parser.add_argument('--height', default=320, type=int, help='Input height') parser.add_argument('--height', default=320, type=int, help='Input height')
parser.add_argument('--width', default=448, type=int, help='Input width') parser.add_argument('--width', default=448, type=int, help='Input width')
parser.add_argument('--proto', '-p', default='FlowNet2_deploy_anysize.prototxt', help='Path to prototxt.') parser.add_argument('--proto', '-p', default='', help='Path to prototxt.')
parser.add_argument('--model', '-m', default='FlowNet2_weights.caffemodel', help='Path to caffemodel.') parser.add_argument('--model', '-m', required=True, help='Path to model.')
args, _ = parser.parse_known_args() args, _ = parser.parse_known_args()
if not os.path.isfile(args.model) or not os.path.isfile(args.proto): if not os.path.isfile(args.model):
raise OSError("Prototxt or caffemodel not exist") raise OSError("Model does not exist")
if args.proto and not os.path.isfile(args.proto):
raise OSError("Prototxt does not exist")
winName = 'Calculation optical flow in OpenCV' winName = 'Calculation optical flow in OpenCV'
cv.namedWindow(winName, cv.WINDOW_NORMAL) cv.namedWindow(winName, cv.WINDOW_NORMAL)
cap = cv.VideoCapture(args.input if args.input else 0) cap = cv.VideoCapture(args.input if args.input else 0)
hasFrame, first_frame = cap.read() hasFrame, first_frame = cap.read()
divisor = 64. if args.proto:
var = {} divisor = 64.
var['ADAPTED_WIDTH'] = int(np.ceil(args.width/divisor) * divisor) var = {}
var['ADAPTED_HEIGHT'] = int(np.ceil(args.height/divisor) * divisor) var['ADAPTED_WIDTH'] = int(np.ceil(args.width/divisor) * divisor)
var['SCALE_WIDTH'] = args.width / float(var['ADAPTED_WIDTH']) var['ADAPTED_HEIGHT'] = int(np.ceil(args.height/divisor) * divisor)
var['SCALE_HEIGHT'] = args.height / float(var['ADAPTED_HEIGHT']) var['SCALE_WIDTH'] = args.width / float(var['ADAPTED_WIDTH'])
var['SCALE_HEIGHT'] = args.height / float(var['ADAPTED_HEIGHT'])
config = '' config = ''
proto = open(args.proto).readlines() proto = open(args.proto).readlines()
for line in proto: for line in proto:
for key, value in var.items(): for key, value in var.items():
tag = "$%s$" % key tag = "$%s$" % key
line = line.replace(tag, str(value)) line = line.replace(tag, str(value))
config += line config += line
caffemodel = open(args.model, 'rb').read() caffemodel = open(args.model, 'rb').read()
opt_flow = OpticalFlow(caffemodel, var['ADAPTED_HEIGHT'], var['ADAPTED_WIDTH'], bytearray(config.encode()))
else:
opt_flow = OpticalFlow(args.model, 360, 480)
opt_flow = OpticalFlow(bytearray(config.encode()), caffemodel, var['ADAPTED_HEIGHT'], var['ADAPTED_WIDTH'])
while cv.waitKey(1) < 0: while cv.waitKey(1) < 0:
hasFrame, second_frame = cap.read() hasFrame, second_frame = cap.read()
if not hasFrame: if not hasFrame: