compute_figure.py 文件源码

def plot_figure_5(images, rigidity_refined, structure_refined, flow_estimated, flow_init, flow_gt, flow_gt_valid):
    if not os.path.isdir('./results'):
        os.makedirs('results')

    I = img_as_ubyte((images[0]+images[1]+images[2])/3.0)
    io.imsave('./results/01_image.png',I)

    cm_bwr = plt.get_cmap('bwr')
    Irigidity = cm_bwr(rigidity_refined.astype('float32'))
    io.imsave('./results/02_rigidity.png',Irigidity)

    Istructure = structure2image(structure_refined, rigidity_refined)
    io.imsave('./results/03_structure.png',Istructure)

    Iuv_est = flow_viz.computeFlowImage(flow_estimated[0],flow_estimated[1])
    io.imsave('./results/04_flow.png',Iuv_est)

    epe_est = np.sqrt((flow_estimated[0]-flow_gt[0])**2 + (flow_estimated[1]-flow_gt[1])**2)
    epe_init = np.sqrt((flow_init[0]-flow_gt[0])**2 + (flow_init[1]-flow_gt[1])**2)

    #import ipdb; ipdb.set_trace()

    epe_est[flow_gt_valid==0] = 0
    epe_init[flow_gt_valid==0] = 0

    epe_diff = epe_init - epe_est
    epe_green = np.clip(epe_diff, 0, 3)/3.0
    epe_red = np.clip(-epe_diff, 0, 3)/3.0

    Icomparison = np.zeros((rigidity_refined.shape[0],rigidity_refined.shape[1],3))

    print(Icomparison.shape)
    print(epe_green.shape)
    print(epe_red.shape)

    Icomparison[:,:,0] = epe_red
    Icomparison[:,:,1] = epe_green
    Icomparison = img_as_ubyte(Icomparison)
    io.imsave('./results/05_comparison.png',Icomparison)



#
# Supmat figures
#