compute_figure.py 文件源码

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

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

    Iuv_gt = flow_viz.computeFlowImage(flow_gt[0], flow_gt[1])
    io.imsave('./results_supmat/temp/02_gt_flow.png', Iuv_gt)

    cm_bwr = plt.get_cmap('bwr')
    Irigidity = cm_bwr(rigidity_refined.astype('float32'))
    io.imsave('./results_supmat/temp/03_rigidity.png',Irigidity)

    Istructure = structure2image(structure_refined, rigidity_refined)
    io.imsave('./results_supmat/temp/04_structure.png',Istructure)

    Iuv_est = flow_viz.computeFlowImage(flow_estimated[0],flow_estimated[1])
    io.imsave('./results_supmat/temp/05_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))

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