Simple TIFF encoder

tiff_encoder.py

from enum import IntEnum


class FieldType(IntEnum):
    BYTE = 1
    ASCII = 2
    SHORT = 3
    LONG = 4
    RATIONAL = 5


class Tag(IntEnum):
    ImageWidth = 0x100
    ImageLength = 0x101
    BitsPerSample = 0x102
    Compression = 0x103
    PhotometricInterpretation = 0x106
    StripOffsets = 0x111
    SamplesPerPixel = 0x115
    RowsPerStrip = 0x116
    StripByteCounts = 0x117
    XResolution = 0x11A
    YResolution = 0x11B
    ResolutionUnit = 0x128


def to_bytes(obj: int, n: int):
    return obj.to_bytes(n, byteorder="little")


def write_field(f, tag: Tag, dtype: FieldType, count: int, value: int):
    f.write(to_bytes(tag, 2))
    f.write(to_bytes(dtype, 2))
    f.write(to_bytes(count, 4))
    f.write(to_bytes(value, 4))


def write_tiff(f):
    f.write(b"II")
    f.write(to_bytes(42, 2))
    f.write(to_bytes(8, 4))

    img_width = 100
    img_height = 100

    n_fields = 11
    offset = 8 + 2 + n_fields * 12 + 4

    f.write(to_bytes(n_fields, 2))
    write_field(f, Tag.ImageWidth, FieldType.SHORT, 1, img_width)
    write_field(f, Tag.ImageLength, FieldType.SHORT, 1, img_height)
    write_field(f, Tag.BitsPerSample, FieldType.SHORT, 3, offset)
    write_field(f, Tag.Compression, FieldType.SHORT, 1, 1)
    write_field(f, Tag.PhotometricInterpretation, FieldType.SHORT, 1, 2)
    write_field(f, Tag.StripOffsets, FieldType.SHORT, 1, offset + 6 + 8 + 8)
    write_field(f, Tag.SamplesPerPixel, FieldType.SHORT, 1, 3)
    write_field(f, Tag.RowsPerStrip, FieldType.SHORT, 1, img_height)
    write_field(f, Tag.StripByteCounts, FieldType.LONG, 1, img_width * img_height * 3)
    write_field(f, Tag.XResolution, FieldType.RATIONAL, 1, offset + 6)  # cannot fit in 4 bytes
    write_field(f, Tag.YResolution, FieldType.RATIONAL, 1, offset + 6 + 8)
    f.write(to_bytes(0, 4))

    f.write(to_bytes(8, 2))  # BitsPerSample
    f.write(to_bytes(8, 2))
    f.write(to_bytes(8, 2))
    f.write(to_bytes(300, 4))  # XResolution
    f.write(to_bytes(1, 4))
    f.write(to_bytes(300, 4))  # YResolution
    f.write(to_bytes(1, 4))

    for i in range(img_height):
        for j in range(img_width):
            idx = i * img_width + j
            color = 256 * idx // (img_height * img_width)
            f.write(to_bytes((color * 1) % 256, 1))
            f.write(to_bytes((color * 2) % 256, 1))
            f.write(to_bytes((color * 4) % 256, 1))


write_tiff(open("sample.tiff", "wb"))