quick hack for reading MetaMorph STK data

try_metamorph.py

"""
From-scratch MetaMorph Stack (STK) reader.

Currently just a quick hack. Dumps the TIFF tags to stdout and each
plane to its own single-image TIFF file.
"""

import sys
import struct
from contextlib import closing

import numpy as np
from libtiff import TIFF  # only for writing out individual planes


MAP_END = {"II": "<", "MM": ">"}

TYPE_MAP = {
    1: (1, "B"),    # BYTE: 8-bit unsigned integer.
    2: (1, "s"),    # ASCII: 8-bit byte that contains a 7-bit ASCII code.
                    # The last byte must be NUL (binary zero)
    3: (1, "H"),    # SHORT: 16-bit (2-byte) unsigned integer.
    4: (1, "I"),    # LONG: 32-bit (4-byte) unsigned integer.
    5: (2, "I"),    # RATIONAL: two LONGs: the first represents the numerator
                    # of a fraction, the second the denominator.
    6: (1, "b"),    # SBYTE: an 8-bit signed integer.
    7: (1, "c"),    # UNDEFINED: an 8-bit byte that may contain anything,
                    # depending on the definition of the field.
    8: (1, "h"),    # SSHORT: a 16-bit (2-byte) signed integer.
    9: (1, "i"),    # SLONG: a 32-bit (4-byte) signed integer.
    10: (2, "i"),   # SRATIONAL: two SLONGs: the first represents the
                    # numerator of a fraction, the second the denominator.
    11: (1, "f"),   # FLOAT: single precision (4-byte) IEEE format.
    12: (1, "d"),   # DOUBLE: double precision (8-byte) IEEE format.
}

# Metamorph tags (special conventions on count)
UIC1 = 33628
UIC2 = 33629
UIC3 = 33630
UIC4 = 33631
MM_TAGS = frozenset((UIC1, UIC2, UIC3, UIC4))


# TODO: replace asserts with proper exceptions; define constants
class TiffTagReader(object):

    def __init__(self, f):
        self.f = f
        self.is_stk = True
        ifd_offset = self.__parse_header()
        print "endianness: %s" % ("BIG" if self.end == ">" else "LITTLE")
        print "first IFD offset: %d" % ifd_offset
        fields, next_offset = self.__parse_ifd(ifd_offset)
        print "N. fields:", len(fields)
        print "next offset:", next_offset
        dump_len = 3
        for tag, value in fields:
            if len(value) <= dump_len:
                dump = repr(value)
            else:
                dump = "%s ...)" % repr(value[:dump_len])[:-1]
            print "  %s (%d) %s" % (tag, len(value), dump)
        # --
        if self.is_stk:
            field_map = dict(fields)
            rows_per_strip = field_map[278][0]
            size_y = field_map[257][0]
            strips_per_image = size_y // rows_per_strip
            strip_offsets = field_map[273]
            strip_byte_counts = field_map[279]
            plane_offsets = []
            for i in xrange(self.n_planes):
                plane_offsets.append(i * (
                    strip_offsets[strips_per_image - 1] +
                    strip_byte_counts[strips_per_image - 1] -
                    strip_offsets[0]
                ))
            print "plane offsets: %s ...]" % \
                repr(plane_offsets[:dump_len])[:-1]
            # --
            print "writing out planes..."
            size_x = field_map[256][0]
            bytes_per_pixel = field_map[258][0] // 8
            plane_size = size_x * size_y * bytes_per_pixel
            for i, offset in enumerate(plane_offsets):
                fname = "plane_%d.tif" % i
                f.seek(plane_offsets[i] + strip_offsets[0])
                plane = f.read(plane_size)
                dtype = np.dtype("%si%d" % (self.end, bytes_per_pixel))
                a = np.fromstring(plane, dtype=dtype).reshape((size_y, size_x))
                with closing(TIFF.open(fname, mode="w")) as fo:
                    fo.write_image(a)

    def __unpack(self, fmt, data):
        return struct.unpack(self.end + fmt, data)

    def __parse_header(self):
        header = self.f.read(8)
        assert len(header) == 8
        self.end = MAP_END.get(header[:2])
        assert self.end is not None
        version, ifd_offset = self.__unpack("HI", header[2:])
        assert version == 42
        return ifd_offset

    def __parse_ifd(self, offset):
        self.f.seek(offset)
        data = self.f.read(2)
        assert len(data) == 2
        n_fields = self.__unpack("H", data)[0]
        size = n_fields * 12 + 4
        data = self.f.read(size)
        fields = [self.__parse_field(data[_: _ + 12])
                  for _ in xrange(0, size - 4, 12)]
        next_offset = self.__unpack("I", data[-4:])[0]
        return fields, next_offset

    def __parse_field(self, data):
        assert len(data) == 12
        tag, field_type, count = self.__unpack("HHI", data[:8])
        # if tag in MM_TAGS:
        if tag in (UIC2, UIC3):  # todo: parse UIC{1,4}
            return self.__handle_mm(data, tag, field_type, count)
        mult, fmt = TYPE_MAP[field_type]
        fmt = "%d%s" % (count * mult, fmt)
        value = self.__read_value(fmt, data)
        if fmt[-1] == "s":
            value = value[0].rstrip("\x00").split("\x00")
        return tag, value

    def __read_value(self, fmt, data):
        size = struct.calcsize(fmt)
        if size <= 4:
            value = self.__unpack(fmt, data[8:8+size])
        else:
            value_offset = self.__unpack("I", data[8:])[0]
            self.f.seek(value_offset)
            value = self.__unpack(fmt, self.f.read(size))
        return value

    def __handle_mm(self, data, tag, field_type, count):
        if tag == UIC2:
            self.is_stk = True
            self.n_planes = count
            fmt = "%dI" % (6 * count)
            value = self.__read_value(fmt, data)
            value = [value[_: _ + 6] for _ in xrange(0, len(value), 6)]
            # first two items in each subseq represent a rational
            value = [((_[:2]),) + _[2:] for _ in value]
            return tag, value
        if tag == UIC3:
            fmt = "%dI" % (2 * count)
            value = self.__read_value(fmt, data)
            value = [value[_: _ + 2] for _ in xrange(0, len(value), 2)]
            return tag, value


with open(sys.argv[1]) as f:
    reader = TiffTagReader(f)