Python script to encode TASBot pixelart script

colorcmds.py

#!/usr/bin/env python

import argparse
import sys

from PIL import Image

def loadimage(result):
    # assuming pixel values are 0..255.. perhaps I should use the image's getextrema func...
    for y in xrange(0, result.height):
        for x in xrange(0, result.width):
            dp = result.getpixel((x,y))
            if dp[3] != 255:
                continue
            yield ((x, y), (dp[0], dp[1], dp[2]))

def main():
    parser = argparse.ArgumentParser(description="Convert an image file to a tasbot color command script. Transparent pixels will be omitted.")

    parser.add_argument("image", metavar="IMG", nargs=1, help="image file to convert", type=argparse.FileType('rb'))
    #parser.add_argument("-f", "--from", dest="source", help="starting image state to work from", type=argparse.FileType('rb'))
    parser.add_argument("-o", "--output", dest="output", help="file to write command script to. Default is stdout", default="-", type=argparse.FileType('w'))
    parser.add_argument("--offset", nargs=4, metavar=("width", "height", "x", "y"), help="specify where the top-left pixel of the input image should be placed in the canvas. Requires the size of the canvas to be specified, in the width and height arguments", type=int)
    parser.add_argument("--group", nargs='?', const='50', type=int, metavar='N', help='group output by color, emiting N squares per line (if N is not specified, defaults to 50)')
    parser.add_argument('--allow-ranges', dest='ranges', action='store_true', default=False, help='allow ranges in group output')
    
    # todo: add options for column-major or row-major addressing of pixels. for now, assume column-major.
    # todo: sort by color?
    # todo: map color to name
    # todo: multiple pixels on same command (requires bot support)

    args = parser.parse_args()
    
    # allow for the source option to be commented out for now, while we don't support it after adding --offset
    args.source = getattr(args, "source", None)
    
    dst = Image.open(args.image[0])

    if args.offset and args.source:
        print >> sys.stderr, "use of --from with --offset is not currently supported"
        return 1
    
    if dst.mode != "RGBA":
        print >> sys.stderr, "file must be RGBA"
        return 1
    
    src = None
    if args.source:
        src = Image.open(args.source)
        if src.mode != "RGBA":
            print >> sys.stderr, "starting image must be RGBA"
            return 1
        if src.size != dst.size:
            print >> sys.stderr, "starting and destination images must be the same dimensions"
            return 1
            
    canvas = dst.size
    if args.offset:
        canvas = (args.offset[0], args.offset[1])
    
    pixels = {}
    outside_count = 0
    for p in loadimage(dst):
        x = p[0][0]
        y = p[0][1]
        if args.offset:
            x += args.offset[2]
            y += args.offset[3]
            if x < 0 or x >= canvas[0] or y < 0 or y >= canvas[1]:
                outside_count += 1
                continue
        loc = x * canvas[1] + y
        color = '#{:02x}{:02x}{:02x}'.format(p[1][0], p[1][1], p[1][2])
        if args.group:
            if color in pixels:
                pixels[color].append(loc)
            else:
                pixels[color] = [loc]
        else:
            pixels[loc] = color
    if outside_count > 0:
        print >> sys.stderr, "Warning: {} pixels were outside the canvas and skipped".format(outside_count)
    
    if args.group:
        def seghelp(grp):
            group = []
            segment = None
            for loc in sorted(grp):
                if not segment:
                    segment = [loc]
                elif args.ranges and loc % canvas[1] > 0 and segment[-1] == loc-1:
                    segment.append(loc)
                else:
                    if len(segment) > 1:
                        group.append('{}-{}'.format(segment[0], segment[-1]))
                    else:
                        group.append(str(segment[0]))
                    segment = [loc]
                    
                    if len(group) >= args.group:
                        yield ','.join(group)
                        group = []
            if len(segment) > 1:
                group.append('{}-{}'.format(segment[0], segment[-1]))
            else:
                group.append(str(segment[0]))
            yield ','.join(group)
            
        color = None
        for clrgrp in sorted(pixels):
            color = clrgrp
            for ln in seghelp(pixels[clrgrp]):
                args.output.write('{} {}\n'.format(color, ln))
    else:
        for loc in sorted(pixels):
            args.output.write('{} {}\n'.format(color, loc))

if __name__ == "__main__":
    main()