How to create a png in Python and control what each pixel is?

Question

I'm trying to create a png in Python but my current code has an issue, the image it makes has to be grayscale.

def makeGrayPNG(data, height = None, width = None):
    def I1(value):
        return struct.pack("!B", value & (2**8-1))
    def I4(value):
        return struct.pack("!I", value & (2**32-1))
    # compute width&height from data if not explicit
    if height is None:
        height = len(data) # rows
    if width is None:
        width = 0
        for row in data:
            if width < len(row):
                width = len(row)
    # generate these chunks depending on image type
    makeIHDR = True
    makeIDAT = True
    makeIEND = True
    png = b"\x89" + "PNG\r\n\x1A\n".encode('ascii')
    if makeIHDR:
        colortype = 0 # true gray image (no palette)
        bitdepth = 8 # with one byte per pixel (0..255)
        compression = 0 # zlib (no choice here)
        filtertype = 0 # adaptive (each scanline seperately)
        interlaced = 0 # no
        IHDR = I4(width) + I4(height) + I1(bitdepth)
        IHDR += I1(colortype) + I1(compression)
        IHDR += I1(filtertype) + I1(interlaced)
        block = "IHDR".encode('ascii') + IHDR
        png += I4(len(IHDR)) + block + I4(zlib.crc32(block))
    if makeIDAT:
        raw = b""
        for y in xrange(height):
            raw += b"\0" # no filter for this scanline
            for x in xrange(width):
                c = b"\0" # default black pixel
                if y < len(data) and x < len(data[y]):
                    c = I1(data[y][x])
                raw += c
        compressor = zlib.compressobj()
        compressed = compressor.compress(raw)
        compressed += compressor.flush() #!!
        block = "IDAT".encode('ascii') + compressed
        png += I4(len(compressed)) + block + I4(zlib.crc32(block))
    if makeIEND:
        block = "IEND".encode('ascii')
        png += I4(0) + block + I4(zlib.crc32(block))
    return png

img = makeGrayPNG([[0,0,0]],500, 1200)

Is there a better way to do this with an easier way to control each pixel as well, I would prefer it if you didn't have to install any modules but I don't mind.

Answer 1

You can work on a numpy array and turn it into a (grayscale) Pillow image, which you can then save as png.

import numpy as np
import PIL
from PIL import Image
rand_np = np.random.rand(200,200)*255  # Creates random 2D array; here goes your grayscale image content
gray_image = Image.fromarray(rand_np).convert('L')
gray_image.save("img1.png")

Edit:

Extending the above to also work with RGB, and following this post, you can do the following:

import numpy as np
import cv2
rgb_content = np.random.rand(200,200,3)*255  # Creates random 2D array; here goes your rgb image content
cv2.imwrite('rgb_img.png', rgb_content, [cv2.IMWRITE_PNG_COMPRESSION, 0])

According to the author of the post referenced above, this is lossless due to [cv2.IMWRITE_PNG_COMPRESSION, 0].