pfm.py

import numpy as np
import re
import sys
import matplotlib.pyplot as plt

def flip(m, axis):
    if not hasattr(m, 'ndim'):
        m = np.asarray(m)
    indexer = [slice(None)] * m.ndim
    try:
        indexer[axis] = slice(None, None, -1)
    except IndexError:
        raise ValueError("axis=%i is invalid for the %i-dimensional input array"
                         % (axis, m.ndim))
    return m[tuple(indexer)]

def flipY(image):
    return flip(image, 0)

def concatChannels(r, g, b):
    assert r.ndim == 2 and g.ndim == 2 and b.ndim == 2
    rgb = (r[..., np.newaxis], g[..., np.newaxis], b[..., np.newaxis])
    return np.concatenate(rgb, axis=-1)

def read(filename):
    try:
        file = open(filename, "r", encoding='iso-8859-1', newline='\n')
    except IOError:
        raise Exception("Can't access 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:
        file.close()
        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:
        file.close()
        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)
    file.close()
    return flipY(np.reshape(data, shape))
	#return np.reshape(data, shape))

def write(filePath, image, scale = 1):
    image = flipY(image)
    file = open(filePath, "w", encoding='iso-8859-1', newline='\n')
    color = None

    if image.dtype.name != 'float32':
        image = np.float32(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:
        file.close()
        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')
    file.write('%d %d\n' % (image.shape[1], image.shape[0]))

    endian = image.dtype.byteorder

    if endian == '<' or endian == '=' and sys.byteorder == 'little':
        scale = -scale

    file.write('%f\n' % scale)

    image.tofile(file)
    file.close()

def combineAndWrite(filePath, imageR, imageG, imageB, scale = 1):
    return write(filePath, concatChannels(imageR, imageG, imageB), scale)

def show(image):
    img = np.clip(image, 0, 1)
    img = np.power(img, 1 / 2.2)
    plt.imshow(img)
    plt.show()