Naive QOI Image Decoder

qoi.py

#!/bin/python3

import struct
from PIL import Image
from sys import argv, exit

HEADER_FORMAT = ">4sIIBB"
HEADER_SIZE = 14

TAG_FORMAT = ">B"
TAG_SIZE = 1

RGB_FORMAT = ">BBB"
RGB_SIZE = 3

RGBA_FORMAT = ">BBBB"
RGBA_SIZE = 4

LUMA_FORMAT = ">B"
LUMA_SIZE = 1

QOI_OP_RGB = 0b11111110
QOI_OP_RGBA = 0b11111111

QOI_OP_INDEX = 0b00
QOI_OP_DIFF = 0b01
QOI_OP_LUMA = 0b10
QOI_OP_RUN = 0b11

def get_mode(channels):
    if channels == 3:
        return "RGB"

    if channels == 4:
        return "RGBA"

    raise ValueError("invalid channels value")

def get_index_position(r, g, b, a):
    return (r * 3 + g * 5 + b * 7 + a * 11) % 64

def parse_qoi(raw):
    signature, width, height, channels, colorspace = struct.unpack(
            HEADER_FORMAT, raw[:HEADER_SIZE])
    if signature != b"qoif":
        raise ValueError("invalid file signature")

    image = Image.new(get_mode(channels), (width, height))

    idx = HEADER_SIZE
    previous = (0, 0, 0, 255)
    array = [(0, 0, 0, 0)] * 64
    pixel = 0

    while pixel < width * height:
        raw_byte, = struct.unpack(TAG_FORMAT, raw[idx:idx+TAG_SIZE])
        idx = idx + TAG_SIZE

        long_tag = raw_byte
        short_tag = raw_byte >> 6

        if long_tag == QOI_OP_RGB:
            r, g, b = struct.unpack(RGB_FORMAT, raw[idx:idx+RGB_SIZE])
            idx = idx + RGB_SIZE
            previous = (r, g, b, 255)
        elif long_tag == QOI_OP_RGBA:
            r, g, b, a = struct.unpack(RGBA_FORMAT, raw[idx:idx+RGBA_SIZE])
            idx = idx + RGBA_SIZE
            previous = (r, g, b, a)
        elif short_tag == QOI_OP_INDEX:
            index = raw_byte & 0b00111111
            previous = array[index]
        elif short_tag == QOI_OP_DIFF:
            dr, dg, db = ((raw_byte >> 4) & 0b0011) - 2, ((raw_byte >> 2) & 0b0011) - 2, (raw_byte & 0b0011) - 2
            r, g, b, a = previous
            r = (r + dr) & 0xff
            g = (g + dg) & 0xff
            b = (b + db) & 0xff
            previous = (r, g, b, a)
        elif short_tag == QOI_OP_LUMA:
            luma, = struct.unpack(LUMA_FORMAT, raw[idx:idx+LUMA_SIZE])
            idx = idx + LUMA_SIZE
            dg = (raw_byte & 0b00111111) - 32
            dr_min_dg, db_min_dg = (luma >> 4) - 8, (luma & 0b1111) - 8
            dr, db = dg + dr_min_dg, dg + db_min_dg
            r, g, b, a = previous
            r = (r + dr) & 0xff
            g = (g + dg) & 0xff
            b = (b + db) & 0xff
            previous = (r, g, b, a)
        elif short_tag == QOI_OP_RUN:
            run = (raw_byte & 0b00111111) + 1
            for i in range(run):
                image.putpixel((pixel % width, pixel // width), previous)
                pixel = pixel + 1
            continue
        else:
            raise ValueError("illegal chunk tag") # impossible

        image.putpixel((pixel % width, pixel // width), previous)
        pixel = pixel + 1
        array[get_index_position(*previous)] = previous

    return image

def main(image_path):
    with open(image_path, "rb") as f:
        raw_image = f.read()

    image = parse_qoi(raw_image)
    image.show()

if __name__ == "__main__":
    if len(argv) < 2:
        print(f"[-] usage: {argv[0]} path")
        exit(1)

    main(argv[1])