OpticalFlowで手ぶれ補正

stabilization.py

import cv2
import numpy


cam = cv2.VideoCapture('in.mp4')

old = cv2.cvtColor(cam.read()[1], cv2.COLOR_BGR2GRAY)

matstack = numpy.identity(3)
matlog = []

while True:
    ok, img = cam.read()
    if not ok:
        break

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    old_features = cv2.goodFeaturesToTrack(old, 400, 0.01, 10)

    cur_features, st, err = cv2.calcOpticalFlowPyrLK(old, gray, old_features, None)

    st[sum(((old_features - cur_features)**2).reshape((len(st), 2)).transpose((1, 0))) > 64**2] = 0

    for old, cur in zip(old_features[st == 1], cur_features[st == 1]):
        old = tuple(old)
        cur = tuple(cur)
 
        cv2.line(img, old, cur, (0, 255, 0), 1)
        cv2.circle(img, cur, 4, (0, 255, 0), 1)

    if st.sum() == 0:
        matlog.append(matstack[:2])
    else:
        mat = cv2.estimateRigidTransform(old_features[st == 1], cur_features[st == 1], False)

        if mat is None:
            matlog.append(matstack[:2])
        else:
            mat = numpy.append(mat, numpy.array([[0, 0, 1]]), axis=0)
            matstack = matstack.dot(mat)
            matlog.append(matstack[:2])

    print(','.join('{0:.10f}'.format(x) for x in matstack[:2].reshape(6)))

    cv2.imshow('frame', img)
    cv2.waitKey(1)

    old = gray

cam.release()

cam = cv2.VideoCapture('in.mp4')
fps = cam.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'X264'), fps, (
    int(cam.get(cv2.CAP_PROP_FRAME_WIDTH)),
    int(cam.get(cv2.CAP_PROP_FRAME_HEIGHT)),
))
out.write(cam.read()[1])

for i, mat in enumerate(matlog):
    ok, img = cam.read()
    if not ok:
        break

    d = matlog[max(0, i - int(fps)):min(i + int(fps), len(matlog) - 1)]
    avg = numpy.append(sum(d) / len(d), numpy.array([[0, 0, 1]]), axis=0)
    mat = numpy.append(mat, numpy.array([[0, 0, 1]]), axis=0)
    dot = numpy.linalg.inv(mat).dot(avg)

    warp = cv2.warpAffine(img, dot[:2], img.shape[1::-1])

    cv2.imshow('smoothing', warp)
    out.write(warp)
    cv2.waitKey(1)

cam.release()
out.release()