Python implementation for converting images to halftone

halftone.py

import cv2
import drawsvg
import argparse
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

def show_raster(im):
    im = cv2.cvtColor(im, cv2.COLOR_GRAY2RGB)
    plt.figure()
    plt.axis("off")
    plt.imshow(im)
    wm = plt.get_current_fig_manager()
    wm.window.showMaximized()
    plt.show()

def scale(size, scale):
    return (size[0]*scale, size[1]*scale)

def main():
    argp = argparse.ArgumentParser(description='Halftone filter.')
    argp.add_argument(
        '--input-file', type=str, required=True,
        help="The name of the file to process.")
    argp.add_argument(
        '--output-file', type=str, 
        help='If empty output is displayed with pyplot.')
    argp.add_argument(
        '--kernel-dim', type=int, default=5,
        help='How many pixels each dot represents. (NxN)')
    argp.add_argument(
        '--dot-radius', type=int, default=None,
        help='The size of each dot.')
    argp.add_argument(
        '--quiet', action='store_true', default=False,
        help='Do not print progress.')
    argp.add_argument(
        '--invert', action='store_true', default=False,
        help='Invert the colors for screen printing.')
    argp.add_argument(
        '--scale', type=int, default=2,
        help='How large of an image to draw. Reduces pixelation.')
    argp.add_argument(
        '--svg', action='store_true', default=False,
        help='Draw an SVG.')
    args = argp.parse_args()

    if args.dot_radius is None:
        args.dot_radius = args.kernel_dim

    bg_color = "black"
    fill_color = "white"

    if args.invert:
        bg_color = "white"
        fill_color = "black"

    if not args.quiet:
        print("Opening image...")
    image = Image.open(args.input_file)

    if not args.quiet:
        print("Converting to grayscale...")
    grayscale_image = image.convert('L')

    pixels = np.array(grayscale_image)
    width, height = grayscale_image.size

    if not args.quiet:
        print("Drawing...")

    if not args.svg:
        halftone_image = Image.new(
            grayscale_image.mode, 
            scale(grayscale_image.size, args.scale))
        draw = ImageDraw.Draw(halftone_image)
        draw.rectangle(
            [0, 0, width*args.scale, height*args.scale], 
            fill=bg_color)
    else:
        draw = drawsvg.Drawing(
            width*args.scale, height*args.scale, 
            origin='top-left')
        draw.append(drawsvg.Rectangle(
            0, 0, width*args.scale, height*args.scale,
            stroke="none",
            fill=bg_color))

    for i in range(0, width, args.kernel_dim):
        offset = (i // args.kernel_dim) % args.kernel_dim
        for j in range(0-offset, height, args.kernel_dim):
            block = pixels[max(j, 0):j+args.kernel_dim, i:i+args.kernel_dim]
            if len(block) == 0:
                continue
            avg = np.mean(block)
            x = (i + block.shape[0] / 2) * args.scale
            y = (j + block.shape[1] / 2) * args.scale
            radius = (args.dot_radius * args.scale / 2) * avg / 255
            if not args.svg:
                draw.ellipse([x - radius, y - radius, 
                    x + radius, y + radius], 
                    fill=fill_color)
            else:
                draw.append(drawsvg.Circle(
                    x, y, radius,
                    stroke="none",
                    fill=fill_color))

    # output_size = scale(grayscale_image.size, args.scale)
    # output_image = halftone_image.resize(output_size)

    if not args.quiet:
        print("Saving...")

    if args.output_file is None:
        if not args.svg:
            show_raster(np.array(halftone_image))
        else:
            raise Exception("SVG not supported for display.")
    else:
        if not args.svg:
            halftone_image.save(args.output_file)
        else:
            draw.save_svg(args.output_file)

if __name__ == "__main__":
    main()