Skip to content

Instantly share code, notes, and snippets.

@liruoteng
Created March 13, 2019 08:58
Show Gist options
  • Save liruoteng/c29f62adeb1a6e365046e239bdc430a5 to your computer and use it in GitHub Desktop.
Save liruoteng/c29f62adeb1a6e365046e239bdc430a5 to your computer and use it in GitHub Desktop.
Create LMDB for FlyingThings3D
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment