add data-processing

Author: DylanWusee

data_preprocess/IO.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3.4
+
+import os
+import re
+import numpy as np
+import uuid
+from scipy import misc
+import numpy as np
+from PIL import Image
+import sys
+
+
+def read(file):
+    if file.endswith('.float3'): return readFloat(file)
+    elif file.endswith('.flo'): return readFlow(file)
+    elif file.endswith('.ppm'): return readImage(file)
+    elif file.endswith('.pgm'): return readImage(file)
+    elif file.endswith('.png'): return readImage(file)
+    elif file.endswith('.jpg'): return readImage(file)
+    elif file.endswith('.pfm'): return readPFM(file)[0]
+    else: raise Exception('don\'t know how to read %s' % file)
+
+def write(file, data):
+    if file.endswith('.float3'): return writeFloat(file, data)
+    elif file.endswith('.flo'): return writeFlow(file, data)
+    elif file.endswith('.ppm'): return writeImage(file, data)
+    elif file.endswith('.pgm'): return writeImage(file, data)
+    elif file.endswith('.png'): return writeImage(file, data)
+    elif file.endswith('.jpg'): return writeImage(file, data)
+    elif file.endswith('.pfm'): return writePFM(file, data)
+    else: raise Exception('don\'t know how to write %s' % file)
+
+def readPFM(file):
+    file = open(file, 'rb')
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header.decode("ascii") == 'PF':
+        color = True
+    elif header.decode("ascii") == 'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii"))
+    if dim_match:
+        width, height = list(map(int, dim_match.groups()))
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().decode("ascii").rstrip())
+    if scale < 0: # little-endian
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>' # big-endian
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data, scale
+
+def writePFM(file, image, scale=1):
+    file = open(file, 'wb')
+
+    color = None
+
+    if image.dtype.name != 'float32':
+        raise Exception('Image dtype must be float32.')
+
+    image = np.flipud(image)
+
+    if len(image.shape) == 3 and image.shape[2] == 3: # color image
+        color = True
+    elif len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1: # greyscale
+        color = False
+    else:
+        raise Exception('Image must have H x W x 3, H x W x 1 or H x W dimensions.')
+
+    file.write('PF\n' if color else 'Pf\n'.encode())
+    file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))
+
+    endian = image.dtype.byteorder
+
+    if endian == '<' or endian == '=' and sys.byteorder == 'little':
+        scale = -scale
+
+    file.write('%f\n'.encode() % scale)
+
+    image.tofile(file)
+
+def readFlow(name):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        return readPFM(name)[0][:,:,0:2]
+
+    f = open(name, 'rb')
+
+    header = f.read(4)
+    if header.decode("utf-8") != 'PIEH':
+        raise Exception('Flow file header does not contain PIEH')
+
+    width = np.fromfile(f, np.int32, 1).squeeze()
+    height = np.fromfile(f, np.int32, 1).squeeze()
+
+    flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2))
+
+    return flow.astype(np.float32)
+
+def readImage(name):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        data = readPFM(name)[0]
+        if len(data.shape)==3:
+            return data[:,:,0:3]
+        else:
+            return data
+
+    return misc.imread(name)
+
+def writeImage(name, data):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        return writePFM(name, data, 1)
+
+    return misc.imsave(name, data)
+
+def writeFlow(name, flow):
+    f = open(name, 'wb')
+    f.write('PIEH'.encode('utf-8'))
+    np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)
+    flow = flow.astype(np.float32)
+    flow.tofile(f)
+
+def readFloat(name):
+    f = open(name, 'rb')
+
+    if(f.readline().decode("utf-8"))  != 'float\n':
+        raise Exception('float file %s did not contain <float> keyword' % name)
+
+    dim = int(f.readline())
+
+    dims = []
+    count = 1
+    for i in range(0, dim):
+        d = int(f.readline())
+        dims.append(d)
+        count *= d
+
+    dims = list(reversed(dims))
+
+    data = np.fromfile(f, np.float32, count).reshape(dims)
+    if dim > 2:
+        data = np.transpose(data, (2, 1, 0))
+        data = np.transpose(data, (1, 0, 2))
+
+    return data
+
+def writeFloat(name, data):
+    f = open(name, 'wb')
+
+    dim=len(data.shape)
+    if dim>3:
+        raise Exception('bad float file dimension: %d' % dim)
+
+    f.write(('float\n').encode('ascii'))
+    f.write(('%d\n' % dim).encode('ascii'))
+
+    if dim == 1:
+        f.write(('%d\n' % data.shape[0]).encode('ascii'))
+    else:
+        f.write(('%d\n' % data.shape[1]).encode('ascii'))
+        f.write(('%d\n' % data.shape[0]).encode('ascii'))
+        for i in range(2, dim):
+            f.write(('%d\n' % data.shape[i]).encode('ascii'))
+
+    data = data.astype(np.float32)
+    if dim==2:
+        data.tofile(f)
+
+    else:
+        np.transpose(data, (2, 0, 1)).tofile(f)
+

data_preprocess/flyingthings3d_utils.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+
+def next_pixel2pc(flow, disparity, save_path=None, f=-1050., cx=479.5, cy=269.5):
+    height, width = disparity.shape
+
+    BASELINE = 1.0
+    depth = -1. * f * BASELINE / disparity
+
+    x = ((np.tile(np.arange(width, dtype=np.float32)[None, :], (height, 1)) - cx + flow[..., 0]) * -1. / disparity)[:,
+        :, None]
+    y = ((np.tile(np.arange(height, dtype=np.float32)[:, None], (1, width)) - cy + flow[..., 1]) * 1. / disparity)[:,
+        :, None]
+    pc = np.concatenate((x, y, depth[:, :, None]), axis=-1)
+
+    if save_path is not None:
+        np.save(save_path, pc)
+    return pc
+
+
+def pixel2pc(disparity, save_path=None, f=-1050., cx=479.5, cy=269.5):
+    height, width = disparity.shape
+
+    BASELINE = 1.0
+    depth = -1. * f * BASELINE / disparity
+
+    x = ((np.tile(np.arange(width, dtype=np.float32)[None, :], (height, 1)) - cx) * -1. / disparity)[:, :, None]
+    y = ((np.tile(np.arange(height, dtype=np.float32)[:, None], (1, width)) - cy) * 1. / disparity)[:, :, None]
+    pc = np.concatenate((x, y, depth[:, :, None]), axis=-1)
+
+    if save_path is not None:
+        np.save(save_path, pc)
+    return pc

data_preprocess/kitti_utils.py

@@ -0,0 +1,62 @@
+import numpy as np
+import png
+
+
+def pixel2xyz(depth, P_rect, px=None, py=None):
+    assert P_rect[0,1] == 0
+    assert P_rect[1,0] == 0
+    assert P_rect[2,0] == 0
+    assert P_rect[2,1] == 0
+    assert P_rect[0,0] == P_rect[1,1]
+    focal_length_pixel = P_rect[0,0]
+    
+    height, width = depth.shape[:2]
+    if px is None:
+        px = np.tile(np.arange(width, dtype=np.float32)[None, :], (height, 1))
+    if py is None:
+        py = np.tile(np.arange(height, dtype=np.float32)[:, None], (1, width))
+    const_x = P_rect[0,2] * depth + P_rect[0,3]
+    const_y = P_rect[1,2] * depth + P_rect[1,3]
+    
+    x = ((px * (depth + P_rect[2,3]) - const_x) / focal_length_pixel) [:, :, None]
+    y = ((py * (depth + P_rect[2,3]) - const_y) / focal_length_pixel) [:, :, None]
+    pc = np.concatenate((x, y, depth[:, :, None]), axis=-1)
+    
+    pc[..., :2] *= -1.
+    return pc
+
+
+def load_uint16PNG(fpath):
+    reader = png.Reader(fpath)
+    pngdata = reader.read()
+    px_array = np.vstack( map(np.uint16, pngdata[2]) )
+    if pngdata[3]['planes'] == 3:
+        width, height = pngdata[:2]
+        px_array = px_array.reshape(height, width, 3)
+    return px_array
+
+
+def load_disp(fpath):
+    # A 0 value indicates an invalid pixel (ie, no
+    # ground truth exists, or the estimation algorithm didn't produce an estimate
+    # for that pixel).
+    array = load_uint16PNG(fpath)
+    valid = array > 0
+    disp = array.astype(np.float32) / 256.0
+    disp[np.logical_not(valid)] = -1.
+    return disp, valid
+
+
+def load_op_flow(fpath):
+    array = load_uint16PNG(fpath)
+    valid = array[..., -1] == 1
+    array = array.astype(np.float32)
+    flow = (array[..., :-1] - 2**15) / 64.
+    return flow, valid
+
+
+def disp_2_depth(disparity, valid_disp, FOCAL_LENGTH_PIXEL):
+    BASELINE = 0.54
+    depth = FOCAL_LENGTH_PIXEL * BASELINE / (disparity + 1e-5)
+    depth[np.logical_not(valid_disp)] = -1.
+    return depth

data_preprocess/process_flyingthings3d_subset.py

@@ -0,0 +1,81 @@
+import numpy as np
+import sys, os
+import os.path as osp
+from multiprocessing import Pool
+import argparse
+
+import IO
+from flyingthings3d_utils import *
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--raw_data_path', type=str, help="path to the raw data")
+parser.add_argument('--save_path', type=str, help="save path")
+parser.add_argument('--only_save_near_pts', dest='save_near', action='store_true',
+                    help='only save near points to save disk space')
+
+args = parser.parse_args()
+root_path = args.raw_data_path
+save_path = args.save_path
+# root_path = osp.abspath('/root/share/data/FlyingThings3D_subset/FlyingThings3D_subset/')
+# truenas_path = osp.abspath('/root/share/data/FlyingThings3D_subset_processed_full')
+# m2t_path = osp.abspath('/m2t/FlyingThings3D_subset_processed_full')
+
+splits = ['train', 'val']
+
+
+def process_one_file(params):
+    try:
+        #import ipdb; ipdb.set_trace()
+        train_val, fname = params
+
+        save_folder_path = osp.join(save_path, train_val, fname)
+        os.makedirs(save_folder_path, exist_ok=True)
+
+        disp1 = IO.read(osp.join(root_path, train_val, 'disparity', 'left', fname + '.pfm'))
+        disp1_occ = IO.read(osp.join(root_path, train_val, 'disparity_occlusions', 'left', fname + '.png'))
+        disp1_change = IO.read(
+            osp.join(root_path, train_val, 'disparity_change', 'left', 'into_future', fname + '.pfm'))
+        flow = IO.read(osp.join(root_path, train_val, 'flow', 'left', 'into_future', fname + '.flo'))
+        flow_occ = IO.read(osp.join(root_path, train_val, 'flow_occlusions', 'left', 'into_future', fname + '.png'))
+
+        pc1 = pixel2pc(disp1)
+        pc2 = next_pixel2pc(flow, disp1 + disp1_change)
+
+        if pc1[..., -1].max() > 0 or pc2[..., -1].max() > 0:
+            print('z > 0', train_val, fname, pc1[..., -1].max(), pc1[..., -1].min(), pc2[..., -1].max(),
+                  pc2[..., -1].min())
+
+        valid_mask = np.logical_and(disp1_occ == 0, flow_occ == 0)
+
+        pc1 = pc1[valid_mask]
+        pc2 = pc2[valid_mask]
+
+        if not args.save_near:
+            np.save(osp.join(save_folder_path, 'pc1.npy'), pc1)
+            np.save(osp.join(save_folder_path, 'pc2.npy'), pc2)
+        else:
+            near_mask = np.logical_and(pc1[..., -1] > -35., pc2[..., -1] > -35.)
+            np.save(osp.join(save_folder_path, 'pc1.npy'), pc1[near_mask])
+            np.save(osp.join(save_folder_path, 'pc2.npy'), pc2[near_mask])
+
+    except Exception as ex:
+        print('error in addressing params', params, 'see exception:')
+        print(ex)
+        sys.stdout.flush()
+        return
+
+
+if __name__ == '__main__':
+    param_list = []
+    for train_val in splits:
+        tmp_path = osp.join(root_path, train_val, 'disparity_change', 'left', 'into_future')
+        param_list.extend([(train_val, item.split('.')[0]) for item in os.listdir(tmp_path)])
+
+    #process_one_file(param_list[0])
+
+    pool = Pool(4)
+    pool.map(process_one_file, param_list)
+    pool.close()
+    pool.join()
+
+    print('Finish all!')