IMG2RBWPPM.py

# -*- coding: utf-8 -*-
"""
Created on Wed Jul 15 08:36:23 2020
@author: lirc572

Convert an image to PPM format for EPD75BHD library (https://github.com/lirc572/EPD75BHD)

This script implements the Floyd–Steinberg dithering algorithm to 
convert images to a format that only has 3 colors: red, black and 
white, without grayscale.

For more information on Floyd–Steinberg algorithm, refer to its Wikipedia page.

Basic usage:
- Make sure you have Pillow installed: python3 -m pip install --upgrade Pillow
- `python3 IMG2RBWPPM.py <path-to-your-image>`
- To use color leveling instead of dithering, uncomment the line `#METHOD = "COLOR_LEVEL"` and comment the line `METHOD = "FLOYD_STEINBERG"`
"""

#METHOD = "COLOR_LEVEL"
METHOD = "FLOYD_STEINBERG"
SHOW = True

import sys
import os

from PIL import Image

def printHelp():
    print('Usage: `python IMG2RGBPPM.py hello.png`')

def euclideanDistance(color1, color2):
    return ((color1[0]-color2[0])**2+(color1[1]-color2[1])**2+(color1[2]-color2[2])**2)**0.5

def findClosestColor(color):
    dist_red   = euclideanDistance(color, (255, 0, 0))
    dist_black = euclideanDistance(color, (0, 0, 0))
    dist_white = euclideanDistance(color, (255, 255, 255))
    if (dist_red<=dist_black and dist_red<=dist_white):
        return (255, 0, 0)
    if (dist_black<=dist_red and dist_black<=dist_white):
        return (0, 0, 0)
    if (dist_white<=dist_black and dist_white<=dist_red):
        return (255, 255, 255)

def colorDiff(color1, color2):
    return (color1[0]-color2[0], color1[1]-color2[1], color1[2]-color2[2])

def colorSum(color1, color2):
    return toValidColor((int(color1[0]+color2[0]), int(color1[1]+color2[1]), int(color1[2]+color2[2])))

def toValidColor(color):
    color = list(color)
    for i in range(3):
        if color[i] < 0:
            color[i] = 0
        elif color[i] > 255:
            color[i] = 255
    return tuple(color)

if __name__ == '__main__':
    if len(sys.argv) not in (2, 3):
        printHelp()
        exit(0)
    src = sys.argv[1]
    tgt = sys.argv[2] if len(sys.argv)==3 else os.path.splitext(src)[0] + '.ppm'
    try:
        img = Image.open(src)
    except:
        print('Error reading image "%s"' % (src,))
    print(img.format, img.size, img.mode)
    if SHOW:
        img.show()
    print('Converting to rbw ppm...')
    if METHOD == "COLOR_LEVEL":
        for y in range(img.size[1]):
            for x in range(img.size[0]):
                old_pixel = img.getpixel((x, y))
                new_pixel = findClosestColor(old_pixel)
                img.putpixel((x, y), new_pixel)
    elif METHOD == "FLOYD_STEINBERG":
        for y in range(img.size[1]):
            for x in range(img.size[0]):
                old_pixel   = img.getpixel((x, y))
                new_pixel   = findClosestColor(old_pixel)
                img.putpixel((x, y), new_pixel)
                quant_error = colorDiff(old_pixel, new_pixel)
                if x+1 < img.size[0]:
                    img.putpixel((x+1, y), colorSum(img.getpixel((x+1, y)), tuple(map(lambda x: x*7/16, quant_error))))
                if x > 0 and y+1 < img.size[1]:
                    img.putpixel((x-1, y+1), colorSum(img.getpixel((x-1, y+1)), tuple(map(lambda x: x*3/16, quant_error))))
                if y+1 < img.size[1]:
                    img.putpixel((x, y+1), colorSum(img.getpixel((x, y+1)), tuple(map(lambda x: x*5/16, quant_error))))
                if x+1 < img.size[0] and y+1 < img.size[1]:
                    img.putpixel((x+1, y+1), colorSum(img.getpixel((x+1, y+1)), tuple(map(lambda x: x/16, quant_error))))
    if SHOW:
        img.show()
    #img.save(tgt)
    ppm_content = "P3\n%d %d\n1\n" % (img.size)
    for y in range(img.size[1]):
        for x in range(img.size[0]):
            px = tuple(img.getpixel((x, y))[:3])
            pix = (1,0,0) if px==(255,0,0) else (1,1,1) if px==(255,255,255) else (0,0,0)
            ppm_content += "%d %d %d\n" % pix
    with open(tgt, 'w') as f:
        f.write(ppm_content)
    print('Conversion finished!')