Manually decoding and viewing PNG file in Python

main.py

import os
import struct
import time
import zlib

import numpy as np
from math import floor
from sfml import sf

png_signature = b'\x89PNG\r\n\x1A\n'

class Image:
    def __init__(self, filename):
        self.name = filename
        self.width, self.height, self.bit_depth, self.color_type, self.compression_method, self.filter_method, \
        self.interlace_method = (0, 0, 0, 0, 0, 0, 0)
        self.background_color = 'None'
        self.color_model = self.compression_name = 'Unknown'

        self.unfiltered = []

        with open(filename, 'rb') as self.file:
            signature, = struct.unpack('!8s', self.file.read(8))

            assert signature == png_signature

            self.chunks = []

            while self.file.tell() != os.fstat(self.file.fileno()).st_size:
                self.chunks.append(self.read_chunk())

            handlers = {
                b'IHDR': self.handle_header,
                b'tTXt': self.handle_text,
                b'zTXt': self.handle_text,
                b'iTXt': self.handle_text,
                b'bKGD': self.handle_background,
                b'IDAT': self.handle_image,
            }

            for chunk in self.chunks:
                print('calling handler for ' + chunk[0].decode('utf-8'))
                handlers.get(chunk[0], self.handle_unknown)(chunk)

    def read_chunk(self):
        chunk_length, chunk_type = struct.unpack('!I4s', self.file.read(8))

        if chunk_length != 0:
            chunk_data = self.file.read(chunk_length)
        else:
            chunk_data = bytes()

        chunk_crc_expected, = struct.unpack('!I', self.file.read(4))

        chunk_crc_calculated = zlib.crc32(chunk_type + chunk_data) & 0xffffffff

        if chunk_crc_expected != chunk_crc_calculated:
            print('wrong crc for chunk "{}"'.format(chunk_type))
            raise

        return chunk_type, chunk_length, chunk_data

    def handle_header(self, chunk):
        _, _, data = chunk

        self.width, self.height, self.bit_depth, self.color_type, self.compression_method, self.filter_method, \
            self.interlace_method = struct.unpack('!IIBBBBB', data)

        color_types = {
            2: 'truecolor (rgb)',
            6: 'truecolor (rgba)'
        }

        if self.compression_method == 0:
            self.compression_name = "deflate"
        else:
            self.compression_name = "illegal"  # TODO

        self.color_model = color_types.get(self.color_type, 'Unknown (' + str(self.color_type) + ')')

    @staticmethod
    def handle_text(chunk):
        name, _, data = chunk

        key = data.split(b'\x00')[0].decode('latin-1', 'ignore')

        method = data[len(key) + 1]

        value = data[len(key) + 2:]
        if name == b'zTXt' or method == 1:
            value = zlib.decompress(value).decode('latin-1', 'ignore')
        else:
            value = value.decode('latin-1', 'ignore')

        # TODO implement iTXt properly
        # keyword
        # NULL
        # compression flag 1 byte
        # compression method 1 byte
        # language tag 0 or more bytes
        # NUL
        # translated keyword 0 or more bytes
        # NUL
        # text 0 or more bytes (rest of data)

        assert method == 0x00

        if name == 'iTXt':
            print(key + ': ' + repr(value))  # remove once implemented properly
        else:
            print(key + ': ' + value)

    def handle_background(self, chunk):
        _, _, data = chunk
        # http://www.libpng.org/pub/png/book/chapter08.html
        # https://www.w3.org/TR/PNG-Decoders.html#D.Background-color
        r, g, b = struct.unpack("!HHH", data)
        self.background_color = '#{:02x}{:02x}{:02x}'.format(r, g, b)

    @staticmethod
    def handle_unknown(chunk):
        name, _, _ = chunk
        print('<========= ' + name.decode('utf-8') + ' =========>')

    @staticmethod
    def timing(f):
        def wrap(*args):
            time1 = time.time()
            ret = f(*args)
            time2 = time.time()
            print('%s function took %0.3f ms' % (f.func_name, (time2 - time1) * 1000.0))
            return ret

        return wrap

    def handle_image(self, chunk):
        _, length, data = chunk

        time1 = time.time()
        decompressed = zlib.decompress(data)
        time2 = time.time()
        print('%s function took %0.3f ms' % ('compressing', (time2 - time1) * 1000.0))


        time1 = time.time()
        scanlines = [
            x.tobytes() for x in np.array_split(
                np.frombuffer(bytearray(decompressed), dtype=np.uint8),
                self.height)
        ]
        time2 = time.time()
        print('%s function took %0.3f ms' % ('scanlines', (time2 - time1) * 1000.0))

        self.unfiltered.append(bytearray(int(len(decompressed) / self.height)))

        prev_scanline = self.unfiltered[0]

        time1 = time.time()
        for filtered in scanlines:
            filter_type = ord(filtered[:1])
            scanline = bytearray(filtered[1:])

            filter_types = {
                0: 'none',
                1: 'sub',
                2: 'up',
                3: 'average',
                4: 'paeth',
            }

            bpp = 3

            def paeth_predictor(a, b, c):
                p = a + b - c
                pa = abs(p - a)
                pb = abs(p - b)
                pc = abs(p - c)
                if pa <= pb and pa <= pc:
                    return a
                elif pb <= pc:
                    return b
                else:
                    return c

            def unknown_unfilter(x):
                print('unhandled scan line of type ' + filter_types.get(filter_type, 'Unknown'))
                raise

            def raw(x): return 0 if x < 0 else scanline[x]

            def prior(x): return 0 if x < 0 else prev_scanline[x]

            unfilter = {
                # none
                0: lambda x: x,
                # sub
                1: lambda x: raw(x) + raw(x - bpp),
                # up
                2: lambda x: raw(x) + prior(x),
                # average
                3: lambda x: raw(x) + floor((raw(x - bpp) + prior(x)) / 2),
                # paeth
                4: lambda x: raw(x) + paeth_predictor(raw(x - bpp), prior(x), prior(x - bpp))
            }

            for x in range(0, len(scanline)):
                scanline[x] = unfilter.get(filter_type, unknown_unfilter)(x) & 0xff

            prev_scanline = scanline
            self.unfiltered.append(prev_scanline)

        time2 = time.time()
        print('%s function took %0.3f ms' % ('unfiltering', (time2 - time1) * 1000.0))
        self.unfiltered = self.unfiltered[1:]
        self.unfiltered = b''.join(x for x in self.unfiltered)

        image_data = bytearray()

        for i, v in enumerate(self.unfiltered):
            image_data.append(v)
            if (i + 1) % 3 == 0:
                image_data.append(0xff)

        self.unfiltered = image_data
        # print(image_data)


def main():
    # image = Image('wallpaper.png')
    # image = Image('3x3-ff7f01ff.png')
    image = Image('wallpaper.png')

    window = sf.RenderWindow(sf.VideoMode(600, 600), "png decoder")

    font = sf.Font.from_file("C:\\Windows\\Fonts\\verdana.ttf")
    text = sf.Text(
        '{} ({}x{})\ncolor model: {}\ndesired background: {}\nbit depth: {}\ncompression: {}\nfilter method: {}\ninterlace method: {}'.
            format(image.name, image.width, image.height, image.color_model, image.background_color, image.bit_depth,
                   image.compression_name, image.filter_method, image.interlace_method), font, 12)
    text.position = sf.Vector2(10, 10)
    text2 = sf.Text("Name: " + image.name, font, 12)
    text2.position = sf.Vector2(10, 10)

    print(len(image.unfiltered))
    print(image.height * 4 * image.width)

    img = sf.Image.from_pixels(image.width, image.height, image.unfiltered)
    tex = sf.Texture.from_image(img)
    sprite = sf.Sprite(tex)

    sprite.position = sf.Vector2(0, 0)

    color = 0
    counter = 0

    while window.is_open:
        counter += 1
        if counter % 100 == 0:
            color = (color + 1) & 0xff
        for event in window.events:
            if event == sf.Event.CLOSED:
                window.close()
            if event == sf.Event.RESIZED:
                window.view = sf.View(sf.Rect(sf.Vector2(0, 0), window.size))

        window.clear(sf.Color(155, 155, 155))
        window.draw(sprite)
        window.draw(text)
        window.display()


main()