Create LMDB for FlyingThings3D

create_flyingthings_lmdb.py

import os
import lmdb
import numpy as np
import pfm
from flowlib import *
from PIL import Image

def read_flow(filename):
    """
    function: read flow file function
    :param flow_path: input flow path ground truth
    :return: flow uv, 2-channel array
    """
    if filename.endswith('.flo'):
        flow = read_flo_file(filename)
    elif filename.endswith('.png'):
        flow = read_png_file(filename)
    elif filename.endswith('.pfm'):
        flow = read_pfm_file(filename)
    else:
        raise Exception('Invalid flow file format!')

    return flow.astype(np.float32)

def generate_new_seq(filename):
    f = open(filename, 'r')
    lines = f.readlines()
    file_list = []
    for line in lines:
        file_list.append(line.strip())
    f.close()
    return file_list

def read_dept(filename):
    """
    Yo, read filename
    :param filename: input depth image path
    :return: depth map, 1-channel 2D array
    """
    if filename.endswith('.pfm'):
        depth_map, _ = pfm.readPFM(filename)
    elif filename.endswith('.png'):  # e.g. sintel
        depth_map = read_image(filename)
    else:
        raise Exception('Invalid depth file format')

    return depth_map.astype(np.float32)
  
def create_crop_lmdb(db_filepath, filelist_root, mode, H, W):
    """
    param: db_filepath: the lmdb files you want to create
    param: filelist_root: at this dir, you prepare a few .txt files containing corresponding images/flow paths
    similar to FlowNet caffe settings
    param: mode='train' or 'test'
    param: H: image height
    param: W: image width
    """
    newH = 384 # new size you want to crop to
    newW = 640 # new size you want to crop to
    # create image pair/flow/disparity pair list
    img1_list = generate_new_seq(os.path.join(filelist_root, (mode+'_img1.txt')))
    img2_list = generate_new_seq(os.path.join(filelist_root, (mode+'_img2.txt')))
    flow_list = generate_new_seq(os.path.join(filelist_root, (mode+'_flow.txt')))
    disp1_list = generate_new_seq(os.path.join(filelist_root, (mode+'_disp1.txt')))
    disp2_list = generate_new_seq(os.path.join(filelist_root, (mode+'_disp2.txt')))
    print("writing LMDB: %s" % db_filepath)
    env = lmdb.open(db_filepath, map_size=int(1e11), writemap=True)
    txn = env.begin(write=True)
    for i in range(len(img1_list)):
        posx, posy = np.random.randint(0, H - newH), np.random.randint(0, H - newW)
        im1 = np.array(Image.open(img1_list[i]))
        im2 = np.array(Image.open(img2_list[i]))
        flow = read_flow(flow_list[i])
        flow = flow[:,:,0:2]
        dp1 = read_dept(disp1_list[i])
        dp2 = read_dept(disp2_list[i])
        # re
        im1 = im1[posx:posx + newH, posy:posy + newW, :]
        im2 = im2[posx:posx + newH, posy:posy + newW, :]
        dp1 = dp1[posx:posx + newH, posy:posy + newW]
        dp2 = dp2[posx:posx + newH, posy:posy + newW]
        flow = flow[posx:posx + newH, posy:posy + newW, :]

        txn.put(img1_list[i], np.reshape(im1, (newH * newW * 3)))
        txn.put(img2_list[i], np.reshape(im2, (newH * newW * 3)))
        txn.put(flow_list[i], np.reshape(flow, (newH * newW * 2)))
        txn.put(disp1_list[i], np.reshape(dp1, (newH * newW)))
        txn.put(disp2_list[i], np.reshape(dp2, (newH * newW)))

    txn.commit()



if __name__ == '__main__':
    create_crop_lmdb('/home/liruoteng/data/flyingthings_train.db',
                '/home/liruoteng/data/FlyingThings/filelists/',
                'train', 540, 960)