-
-
Save mizar/4330cbe3c8bd4917b1d149c072907a17 to your computer and use it in GitHub Desktop.
Generating a PNG in Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import base64, io, struct, zlib | |
from typing import BinaryIO, List, Tuple, Union | |
COLOR_TYPE_V = 0 | |
COLOR_TYPE_VA = 4 | |
COLOR_TYPE_RGB = 2 | |
COLOR_TYPE_RGBA = 6 | |
PixelV = Tuple[int] | |
PixelVA = Tuple[int, int] | |
PixelRGB = Tuple[int, int, int] | |
PixelRGBA = Tuple[int, int, int, int] | |
ImageV = Tuple[int, List[List[PixelV]]] | |
ImageVA = Tuple[int, List[List[PixelVA]]] | |
ImageRGB = Tuple[int, List[List[PixelRGB]]] | |
ImageRGBA = Tuple[int, List[List[PixelRGBA]]] | |
def generate_checkerboard_pattern(width: int, height: int, color_type = COLOR_TYPE_RGB) -> Union[ImageV, ImageVA, ImageRGB, ImageRGBA]: | |
black_pixel, white_pixel = { | |
COLOR_TYPE_V: ((0,), (255,)), | |
COLOR_TYPE_VA: ((0, 255,), (255, 255,)), | |
COLOR_TYPE_RGB: ((0, 0, 0,), (255, 255, 255,)), | |
COLOR_TYPE_RGBA: ((0, 0, 0, 255,), (255, 255, 255, 255,)), | |
}[color_type] | |
pixels = [] | |
for i in range(height): | |
# Generate a single row of white/black pixels | |
row = [] | |
for j in range(width): | |
if (i + j) % 2 == 0: | |
row.append(white_pixel) | |
else: | |
row.append(black_pixel) | |
pixels.append(row) | |
return (color_type, pixels) | |
def get_checksum(chunk_type: bytes, data: bytes) -> int: | |
checksum = zlib.crc32(chunk_type) | |
checksum = zlib.crc32(data, checksum) | |
return checksum | |
def chunk(out: BinaryIO, chunk_type: bytes, data: bytes) -> None: | |
out.write(struct.pack('>I', len(data))) | |
out.write(chunk_type) | |
out.write(data) | |
checksum = get_checksum(chunk_type, data) | |
out.write(struct.pack('>I', checksum)) | |
def make_ihdr(width: int, height: int, bit_depth: int, color_type: int) -> bytes: | |
return struct.pack('>2I5B', width, height, bit_depth, color_type, 0, 0, 0) | |
def encode_data(img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> List[int]: | |
ret = [] | |
for row in img[1]: | |
ret.append(0) # Filter Method Type: None | |
color_values = [ | |
color_value | |
for pixel in row | |
for color_value in pixel | |
] | |
ret.extend(color_values) | |
return ret | |
def compress_data(data: List[int]) -> bytes: | |
data_bytes = bytearray(data) | |
return zlib.compress(data_bytes, level=9) | |
def make_idat(img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> bytes: | |
encoded_data = encode_data(img) | |
compressed_data = compress_data(encoded_data) | |
return compressed_data | |
def dump_png(out: BinaryIO, img: Union[ImageV, ImageVA, ImageRGB, ImageRGBA]) -> None: | |
HEADER = b'\x89PNG\x0d\x0a\x1a\x0a' | |
out.write(HEADER) # start by writing the header | |
assert len(img[1]) > 0 # assume we were not given empty image data | |
width = len(img[1][0]) | |
height = len(img[1]) | |
bit_depth = 8 # bits per pixel | |
color_type = img[0] | |
ihdr_data = make_ihdr(width, height, bit_depth, color_type) | |
chunk(out, b'IHDR', ihdr_data) | |
compressed_data = make_idat(img) | |
chunk(out, b'IDAT', compressed_data) | |
chunk(out, b'IEND', b'') | |
if __name__ == '__main__': | |
width = 64 | |
height = 64 | |
for color_type, fname in [(COLOR_TYPE_V, 'out_v.png'), (COLOR_TYPE_VA, 'out_va.png'), (COLOR_TYPE_RGB, 'out_rgb.png'), (COLOR_TYPE_RGBA, 'out_rgba.png')]: | |
with io.BytesIO() as membuf: | |
dump_png(membuf, generate_checkerboard_pattern(width, height, color_type)) | |
with open(fname, 'wb') as out: | |
out.write(membuf.getvalue()) | |
print((b'data:image/png;base64,' + base64.b64encode(membuf.getvalue())).decode('utf-8')) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment