update readme and add source code

deformpiv.py

@@ -0,0 +1,276 @@
+import os
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+from openpiv import tools, scaling, pyprocess, validation, filters
+from UnFlowNet.models import Network, device, estimate
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+""" Basic functions
+"""
+
+class AttrDict(dict):
+    __setattr__ = dict.__setitem__
+    __getattr__ = dict.__getitem__
+
+
+# interpolation with opencv
+def remap(img, x, y):
+    x, y = np.float32(x), np.float32(y)
+    out = cv2.remap(img, x, y, cv2.INTER_CUBIC)  # INTER_LANCZOS4 INTER_CUBIC INTER_LINEAR
+    return out
+
+
+# griddes sparse vector to dense field
+def sparse2dense(x, y, u, v, sz):
+    dx, dy = np.meshgrid(np.arange(sz[1]), np.arange(sz[0]))
+    dx = (dx-x[0,0])/(x[0,1]-x[0,0])
+    dy = (dy-y[0,0])/(y[1,0]-y[0,0])
+    du = remap(u, dx, dy)
+    dv = remap(v, dx, dy)
+    return dx, dy, du, dv
+
+
+# piv kernel with Open PIV package
+def openpiv(frame_a, frame_b, winsz=32, overlap=24):
+    # process image pair with extended search area piv algorithm.
+    u, v, sig2noise = pyprocess.extended_search_area_piv( frame_a, frame_b, \
+        window_size=winsz, overlap=overlap, dt=1.0, search_area_size=winsz, sig2noise_method='peak2peak')
+    u1, v1, mask = validation.sig2noise_val(u.copy(), v.copy(), sig2noise, threshold = 1.5 )
+    u = u if np.sum(mask)==np.prod(mask.shape) else u1
+    v = v if np.sum(mask)==np.prod(mask.shape) else v1
+    u, v = filters.replace_outliers( u, v, method='localmean', max_iter=10, kernel_size=2)
+    # get window centers coordinates
+    x, y = pyprocess.get_coordinates( image_size=frame_a.shape, search_area_size=winsz, overlap=overlap)
+    return x, y, u, v
+
+
+# piv kernel with optical flow (opencv)
+def opticalflow(img1, img2, level=4):
+    flow1 = cv2.calcOpticalFlowFarneback(img1, img2, None, 0.5, level, 33, 11, 9, 1.3, 0)
+    flow2 = cv2.calcOpticalFlowFarneback(img2, img1, None, 0.5, level, 33, 11, 9, 1.3, 0)
+    flow = (flow1-flow2)/2
+    u, v = flow[...,0], flow[...,1]
+    x, y = np.meshgrid(np.arange(img1.shape[1]), np.arange(img1.shape[0]))
+    return x, y, u, v
+
+
+# piv kernel with deep neural networks
+""" https://github.com/erizmr/UnLiteFlowNet-PIV
+Unsupervised learning of Particle Image Velocimetry
+"""
+def _loadNet():
+    PATH = './UnFlowNet/UnsupervisedLiteFlowNet_pretrained.pt'
+    unliteflownet.load_state_dict(torch.load(PATH)['model_state_dict'])
+    unliteflownet.eval()
+    unliteflownet.to(device)
+    print('unliteflownet load successfully.')
+
+unliteflownet = Network()
+_loadNet()
+
+def deeppiv1(img1, img2):
+    # The input of the network is recommended to be (256, 256)
+    assert img1.shape == img2.shape #== (256,256)
+    sz = img1.shape
+    h, w = sz[0], sz[1]
+    x1 = torch.Tensor(img1/255.0).view(1,1,sz[0],sz[1])
+    x2 = torch.Tensor(img2/255.0).view(1,1,sz[0],sz[1])
+
+    if img1.shape != (256, 256):
+        x1 = F.interpolate(x1, (256, 256), mode='bilinear', align_corners=False)
+        x2 = F.interpolate(x2, (256, 256), mode='bilinear', align_corners=False)
+
+    y_pre = estimate(x1.to(device), x2.to(device), unliteflownet, train=False)
+    y_pre = F.interpolate(y_pre, (h, w), mode='bilinear', align_corners=False)
+
+    u = y_pre[0][0].detach()
+    v = y_pre[0][1].detach()
+
+    u = u.numpy()
+    v = v.numpy()
+
+    x, y = np.meshgrid(np.arange(img1.shape[1]), np.arange(img1.shape[0]))
+    return x, y, u, v
+
+
+def deeppiv(img1, img2):
+    # The input of the network is recommended to be (256, 256) 
+    assert img1.shape == img2.shape #== (256,256)
+    sz = img1.shape
+    h, w = sz[0], sz[1]
+    x1 = torch.Tensor(img1/255.0).view(1,1,sz[0],sz[1])
+    x2 = torch.Tensor(img2/255.0).view(1,1,sz[0],sz[1])
+
+    if img1.shape[0] != 256 or img1.shape[1] != 256: 
+        s = 16
+
+        # padding zero to 256+n*(256-2s) >= h+2s
+        nx, ny = (h+2*s-256-1)//(256-2*s)+2, (w+2*s-256-1)//(256-2*s)+2
+        px = 256+(nx-1)*(256-2*s)-h
+        py = 256+(ny-1)*(256-2*s)-w
+
+        pad_x1 = F.pad(x1, (s, py-s, s, px-s), "constant", 0)
+        pad_x2 = F.pad(x2, (s, py-s, s, px-s), "constant", 0)
+        hp1, wp1 = pad_x1.shape[2], pad_x1.shape[3] 
+    #     print("padding size", pad_x1.shape)
+
+        # change the shape to B*1*256*256, reorganise the large image to patches
+        win_x1 = nn.Unfold(kernel_size=(256, 256), stride=((256-2*s), (256-2*s)))(pad_x1)
+        win_x1 = win_x1.permute(2, 0, 1)  # B*N*(w*h)
+        win_x1 = win_x1.reshape(win_x1.shape[0], win_x1.shape[1], 256, 256)  # B*N*w*h
+        win_x2 = nn.Unfold(kernel_size=(256, 256), stride=((256-2*s), (256-2*s)))(pad_x2)
+        win_x2 = win_x2.permute(2, 0, 1)  # B*N*(w*h)
+        win_x2 = win_x2.reshape(win_x2.shape[0], win_x2.shape[1], 256, 256)  # B*N*w*h
+    #     print("win_x1.shape", win_x1.shape, x1.shape)
+
+        with torch.no_grad():
+            torch.cuda.empty_cache()
+            y_pre = estimate(win_x1.to(device), win_x2.to(device), unliteflownet, train=False)
+
+        y_pre = y_pre[:, :, s:-s, s:-s]
+        # change back
+        sz = (nx, ny, 2, y_pre.shape[2], y_pre.shape[3])
+        y_pre = y_pre.reshape(sz)
+        y_pre = y_pre.permute(2, 0, 3, 1, 4) # B*2*256*256-> 2*B*256*256
+        y_pre = y_pre.reshape(2, (256-2*s)*nx, (256-2*s)*ny)
+
+        y_pre = y_pre.detach().numpy()
+        u, v = y_pre[0,:,:], y_pre[1,:,:]
+        u = u[:h, :w]
+        v = v[:h, :w]
+        assert u.shape == img1.shape # check the shape
+    else:
+        with torch.no_grad():
+            torch.cuda.empty_cache()
+            y_pre = estimate(x1.to(device), x2.to(device), unliteflownet, train=False)
+        u = y_pre[0][0].detach().numpy()
+        v = y_pre[0][1].detach().numpy()
+
+    x, y = np.meshgrid(np.arange(img1.shape[1]), np.arange(img1.shape[0]))
+    return x, y, u, v
+
+
+""" image warping with deformation field or velocity field.
+"""
+def deform(u, v, delta=1, n_iter=10):
+    """ Generate deformation field
+    Integrates a vector field via scaling and squaring.
+    adopted from https://github.com/voxelmorph/voxelmorph/blob/dev/voxelmorph/torch/layers.py
+    """
+    assert u.shape == v.shape
+
+    x, y = np.meshgrid(np.arange(u.shape[1]), np.arange(u.shape[0]))
+    u, v = u/delta, v/delta
+    dx, dy = u/2**n_iter, v/2**n_iter
+
+    for iter in range(n_iter):
+        dx_new = dx + remap(dx, x+dx, y+dy)
+        dy_new = dy + remap(dy, x+dx, y+dy)
+        dx, dy = dx_new, dy_new
+
+    dx, dy = dx*delta, dy*delta
+    return dx, dy
+
+
+def warping(img1, img2, u, v, method='CDI'):
+    # FDI: out1(x,y)=img1(x+u, y+v)
+    # FDI, CDI, FDDI, CDDI
+
+    assert img1.shape == img2.shape == u.shape == v.shape
+
+    x, y = np.meshgrid(np.arange(u.shape[1]), np.arange(u.shape[0]))
+    u, v = -u, -v
+    if method == 'FDI':
+        out1= remap(img1, x+u, y+v)
+        out2= img2.copy()
+    elif method == 'FDI2':
+        out1= img1.copy()
+        out2= remap(img2, x-u, y-v)
+    elif method == 'CDI':
+        out1= remap(img1, x+0.5*u, y+0.5*v)
+        out2= remap(img2, x-0.5*u, y-0.5*v)
+    elif method == 'FDDI':
+        dx, dy = deform(u, v, delta=1)
+        out1= remap(img1, x+dx, y+dy)
+        out2= img2.copy()
+    elif method == 'FDDI2':
+        dx, dy = deform(-u, -v, delta=1)
+        out1= img1.copy()
+        out2= remap(img2, x+dx, y+dy)
+    elif method == 'CDDI':
+        dx1, dy1 = deform(0.5*u, 0.5*v, delta=1)
+        dx2, dy2 = deform(-0.5*u, -0.5*v, delta=1)
+        out1= remap(img1, x+dx1, y+dy1)
+        out2= remap(img2, x+dx2, y+dy2)
+    else:
+        raise NotImplementedError
+    return out1, out2
+
+
+# Our wrapper for iterative deformation PIV
+class DeformPIV():
+    def __init__(self, config):
+        self._c = config
+        assert self._c.pivmethod in ['opticalflow', 'openpiv', 'deeppiv']
+        assert self._c.deform in ['FDI', 'FDI2', 'CDI', 'FDDI', 'FDDI2', 'CDDI']
+
+        self.onepass = eval(self._c.pivmethod) # opticalflow1
+        self.warping = self._c.deform
+
+    def compute(self, image1, image2, u=None, v=None):
+        # obtain the initial vector field
+        if u is not None:
+            assert image1.shape == image2.shape == u.shape == v.shape
+        else:
+            assert image1.shape == image2.shape
+            x, y, u, v = self.onepass(image1, image2)
+
+        # iterative operation
+        for iter in range(self._c.runs):
+            # using a blur trick to make the iteration stable
+            smooth_k = 3 if u.shape != image1.shape else 9
+            for i in range(2):
+                u = cv2.blur(u, (smooth_k,smooth_k)) # using 19
+                v = cv2.blur(v, (smooth_k,smooth_k))
+
+            # sparse vector to dense field, if needed
+            if u.shape != image1.shape:
+                xd, yd, ud, vd = sparse2dense(x, y, u, v, image1.shape)
+            else:
+                ud, vd = u, v
+
+
+            # image warping
+            # warp = self.warping if iter > -2 else 'CDI'
+            img1, img2 = warping(image1, image2, ud, vd, self.warping)
+
+            # update the estimation
+            if self._c.pivmethod == 'opticalflow':
+                x, y, du, dv = self.onepass(img1, img2, level=0)
+            elif self._c.pivmethod == 'openpiv':
+                # x, y, du, dv = self.onepass(img1, img2, winsz=16, overlap=8)
+                x, y, du, dv = self.onepass(img1, img2)
+            elif self._c.pivmethod == 'deeppiv':
+                x, y, du, dv = self.onepass(img1, img2)
+            else:
+                raise NotImplementedError
+
+            if u.shape !=du.shape:
+                u = remap(ud, x, y)
+                v = remap(ud, x, y)
+
+            u, v = u+du, v+dv
+            # print('The mean of increasing amplitude:',np.mean(np.sqrt(du**2+dv**2)))
+
+        # # Debug plot, save the images as .png files
+        # for k, img in enumerate([image1, img1, img2, image2]):
+        #     fig = plt.figure(figsize=(12,12))
+        #     plt.imshow(img)
+        #     plt.savefig(f"{iter}img{k+1}.png")
+        #     plt.close(fig)
+        return x, y, u, v