PPM to PBM to PNG

pbm_to_png.py

from PIL import Image
from io import BytesIO
import itertools


def as_pbm(w, h, data):  # -> bytes
    ret = bytes()
    ret += b'P4\n'
    ret += (str(w) + ' ' + str(h) + '\n').encode('ascii')
    ret += plain_to_binary(data)
    return ret


def read_pbm(data):
    return ''.join(format(x, '08b') for x in data)


def plain_to_binary(data):  # -> bytes
    lines = (d.replace(' ', '') for d in data.strip().split('\n'))
    ints = itertools.chain.from_iterable(bits_to_bytes(l) for l in lines)
    return b''.join(int_to_bytes(i) for i in ints)


def bits_to_bytes(bits):
    padded_bytes = itertools.zip_longest(*(iter(bits),) * 8, fillvalue='0')
    padded_bytes = ([str.join('', z) for z in padded_bytes])
    return [int(byte, 2) for byte in padded_bytes]


def int_to_bytes(i):
    return i.to_bytes((i.bit_length() + 7) // 8, 'big') or b'\0'


def as_png(w, h, data):  # -> bytes
    pbm_bytes = as_pbm(w, h, data)
    infile = BytesIO(pbm_bytes)
    outfile = BytesIO()
    img = Image.open(infile)
    img.save(outfile, 'PNG')
    return outfile.getvalue()


data = '''
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
0 1 1 1 1 0 0 1 1 1 1 0 0 1 1 1 1 0 0 1 1 1 1 0
0 1 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 1 0 0 1 0
0 1 1 1 0 0 0 1 1 1 0 0 0 1 1 1 0 0 0 1 1 1 1 0
0 1 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0
0 1 0 0 0 0 0 1 1 1 1 0 0 1 1 1 1 0 0 1 0 0 0 0
0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0'''

assert read_pbm(plain_to_binary(data)) == data.replace('\n', '').replace(' ', '')

open('test.png', 'wb').write(as_png(24, 7, data))