Skip to content

Instantly share code, notes, and snippets.

@el3ment
Created May 3, 2017 18:36
Show Gist options
  • Save el3ment/8f22539ad1770423771bf376ddafde58 to your computer and use it in GitHub Desktop.
Save el3ment/8f22539ad1770423771bf376ddafde58 to your computer and use it in GitHub Desktop.
import numpy as np
import re
import sys
from tqdm import tqdm
TAG_CHAR = np.array([202021.25], np.float32)
'''
Load a PFM file into a Numpy array. Note that it will have
a shape of H x W, not W x H. Returns a tuple containing the
loaded image and the scale factor from the file.
'''
def load_pfm(file):
color = None
width = None
height = None
scale = None
endian = None
header = file.readline().rstrip()
if header == 'PF':
color = True
elif header == 'Pf':
color = False
else:
raise Exception('Not a PFM file.')
dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline())
if dim_match:
width, height = map(int, dim_match.groups())
else:
raise Exception('Malformed PFM header.')
scale = float(file.readline().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)
return np.reshape(data, shape), scale
def writeFlow(filename,uv,v=None):
""" Write optical flow to file.
If v is None, uv is assumed to contain both u and v channels,
stacked in depth.
Original code by Deqing Sun, adapted from Daniel Scharstein.
"""
nBands = 2
if v is None:
assert(uv.ndim == 3)
assert(uv.shape[2] == 2)
u = uv[:,:,0]
v = uv[:,:,1]
else:
u = uv
assert(u.shape == v.shape)
height,width = u.shape
f = open(filename,'wb')
# write the header
f.write(TAG_CHAR)
np.array(width).astype(np.int32).tofile(f)
np.array(height).astype(np.int32).tofile(f)
# arrange into matrix form
tmp = np.zeros((height, width*nBands))
tmp[:,np.arange(width)*2] = u
tmp[:,np.arange(width)*2 + 1] = v
tmp.astype(np.float32).tofile(f)
f.close()
if __name__ == '__main__':
import os
for filename in tqdm([os.path.join("./",fn) for fn in next(os.walk("./"))[2]]):
try:
with open(filename) as file:
data, scale = load_pfm(file)
flowFilename = filename.replace('.pfm', '.flo')
writeFlow(flowFilename, data[:,:,0:2])
except Exception as e:
print '{}had error'.format(filename)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment