optical strain（python）

利用python opencv计算optical strain

import cv2
import numpy as np

class OptFlowStrain:
    def __init__(self):
        self.TVL1 = cv2.optflow.DualTVL1OpticalFlow_create()

    def run(self, img0, img1):
        return self.cal_opt_flow(img0, img1)

    def get_opt_res(self):
        return

    def get_translation_matrix(self, nx, ny):
        M = np.float32([[1, 0, nx], [0, 1, ny]])
        return M

    def translation(self, src, nx, ny):
        h, w = src.shape
        # print('h: {}, w: {}'.format(h,w))
        M = self.get_translation_matrix(nx, ny)
        dst = cv2.warpAffine(src, M, (w,h))
        return dst

    def split_opt_flow(self, flow):
        xy0, xy1 = flow[:, :, 0], flow[:, :, 1]
        # cv2.cartToPolar()
        mag, ang = cv2.cartToPolar(xy0, xy1)#, angleInDegrees=True
        u, v = cv2.polarToCart(mag, ang)

        px = self.translation(u, -1, 0)
        py = self.translation(u, 0, 0-1)
        ux = u-px
        uy = u-py

        qx = self.translation(v,-1,0)
        qy = self.translation(v,0,-1)
        vx = v-qx
        vy = v-qy

        temp = cv2.add(uy, vx)
        os = np.multiply(ux, ux) + np.multiply(vy, vy) + 0.5*np.multiply(temp, temp)
        return u, v, cv2.sqrt(os)

    def cv32fc1_to_cv_8uc1(self, source):
        assert source.dtype == np.float32
        source.astype(np.uint8)
        source *= 255
        return source

    def cal_opt_flow(self, img0, img1):
        assert img0.dtype == np.uint8
        assert img1.dtype == np.uint8
        assert img0.shape == img1.shape

        flow = self.TVL1.calc(img0, img1, None)

        u,v,s = self.split_opt_flow(flow)

        u = self.cv32fc1_to_cv_8uc1(u)
        v = self.cv32fc1_to_cv_8uc1(v)
        s = self.cv32fc1_to_cv_8uc1(s)

        # cv2.imwrite('u.jpg', u)
        # cv2.imwrite('v.jpg', v)
        # cv2.imwrite('s.jpg', s)

        return u, v, s