sfm_opticalflow.py

import numpy as np
import cv2
import time,sys,os

ply_header = '''ply
format ascii 1.0
element vertex %(vert_num)d
property float x
property float y
property float z
property uchar red
property uchar green
property uchar blue
end_header
'''

def write_ply(fn, verts, colors):
    verts = verts.reshape(-1, 3)
    colors = colors.reshape(-1, 3)
    verts = np.hstack([verts, colors])
    with open(fn, 'w') as f:
        f.write(ply_header % dict(vert_num=len(verts)))
        np.savetxt(f, verts, '%f %f %f %d %d %d')

def getflow(flow,step=10):
    h,w = flow.shape[:2]
    y,x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1) # (0,0)(0,1)(0,2)(1,0)...
    pts1 = np.vstack((y,x))
    u,v = flow[y,x].T
    y += u
    x += v
    pts2 = np.vstack((y,x))
    return np.float64(pts1).T, np.float64(pts2).T   

imgl = cv2.imread('0001.png',0)
imgr = cv2.imread('0002.png',0)
imgc = cv2.imread('0001.png')
colors = cv2.cvtColor(imgc, cv2.COLOR_BGR2RGB)

######                    step 1 : Find the correspondences           ##################

flow = cv2.calcOpticalFlowFarneback(imgl,imgr, None, 0.5, 3, 15, 3, 5, 1.2, 0)
pts1, pts2 = getflow(flow,5)

# ########################################### step 2: find fundamental matrix, E
F, mask = cv2.findFundamentalMat(pts1,pts2)
pts11 = pts1[np.bool_(mask.ravel())]
pts22 = pts2[np.bool_(mask.ravel())]

# remove those points outside image
r,c = imgc.shape[:2]
valid = (pts22[:,0]<r) & (pts22[:,1]<c)
pts111 = pts11[valid]
pts222 = pts22[valid]

K = np.loadtxt('0000.png.K')
E = K.T.dot(F).dot(K)

R1,R2,t = cv2.decomposeEssentialMat(E)

for i,P in enumerate(((R1,t),(R1,-t),(R2,t),(R2,-t))):

    P1 = K.dot(np.hstack(P))       # Projection matrix of second cam is ready

    P00 = np.float64([ [1,0,0,0],
                    [0,1,0,0],
                    [0,0,1,0]   ]) # Projection matrix of first cam at origin
    P0 = K.dot(P00) 
              
    ########################################## Step 3 : Triangulation

    pts1u = cv2.undistortPoints(pts111.reshape(-1,1,2),K,None)
    pts2u = cv2.undistortPoints(pts222.reshape(-1,1,2),K,None)

    res = cv2.triangulatePoints(P0,P1,pts1u,pts2u)  # (-1,4)
    pts = cv2.convertPointsFromHomogeneous(res.T).reshape(-1,3)

    pts111_int = np.int0(pts111)
    color = imgc[pts111_int[:,0].ravel(), pts111_int[:,1].ravel()]
    
    write_ply('outflow.ply',pts,color)
    os.system('meshlab outflow.ply')