slow floodfill

#!/usr/bin/env python

# slow floodfill for /tg/

from __future__ import print_function
from PIL import Image, GifImagePlugin
import sys
import struct
import subprocess
import optparse

# hard-coded gif color indexes -- change these to suit the image, or vice-versa
#
# floodfilling will start from FILL pixels and fill EMPTY pixels if you don't
# provide a starting coordinate
# otherwise, just set FILL, FILLEDGE, and TRANSPARENT to something unoccupied
WALL=0
EMPTY=1
INVISIBLEWALL=2
FILL=3
FILLEDGE=4
TRANSPARENT=255

def neighbors(pixel):
   x, y = pixel
   return [(x-1,y), (x+1, y), (x, y-1), (x, y+1)]

def moore_neighbors(pixel):
   x, y = pixel
   return [(x-1,y-1), (x, y-1), (x+1, y-1), (x-1, y), (x+1, y), (x-1, y+1), (x, y+1), (x+1, y+1)]

def fillpixel(unfilled, pixel, neighborfn):
   nextfilled = set()
   for neighbor in neighborfn(pixel):
      if neighbor in unfilled:
         nextfilled.add(neighbor)
   return nextfilled

def allneighbors(unfilled, tofill, neighborfn):
   nextfilled = set()
   for pixel in tofill:
      neighbors = fillpixel(unfilled, pixel, neighborfn)
      unfilled.difference_update(neighbors)
      nextfilled.update(neighbors)
   return nextfilled

def fillpixels(img, unfilled, tofill, neighborfn, frames=1):
   nextfilled = set()
   for pixel in tofill:
      try:
         img.putpixel(pixel, FILL)
      except IndexError:
         pass
   for frame in range(frames):
      tofill = allneighbors(unfilled, tofill, neighborfn)
      nextfilled.update(tofill)
   for pixel in nextfilled:
      img.putpixel(pixel, FILLEDGE)
   return nextfilled

def setpalette(img):
   palette = img.getpalette()
   palette[FILLEDGE*3:FILLEDGE*3+3] = [0x00, 0x00, 0xff]
   palette[FILL*3:FILL*3+3] = [0xff, 0x00, 0x00]
   palette[INVISIBLEWALL*3:INVISIBLEWALL*3+3] = palette[WALL*3:WALL*3+3]
   img.putpalette(palette)
   img.info["transparency"] = TRANSPARENT

def getcolorpixels(img):
   colors = {}
   maxx, maxy = img.size
   for y in range(maxy):
      for x in range(maxx):
         pixel = (x,y)
         color = img.getpixel(pixel)
         colors.setdefault(color, set()).add(pixel)
   return colors

def clearimg(img, pixel):
   maxx, maxy = img.size
   img.putdata([pixel for x in range(maxx * maxy)])

def floodfill_frames(img, neighborfn, startcoord=None, framestep=1):
   maxx, maxy = img.size
   emptyframe = img.copy()
   clearimg(emptyframe, TRANSPARENT)
   colors = getcolorpixels(img)
   if startcoord == None:
      unfilled = colors[EMPTY]
      nextfilled = allneighbors(unfilled, colors[FILL], neighborfn)
   else:
      startx = min(maxx, max(0, startcoord[0]))
      starty = min(maxy, max(0, startcoord[1]))
      startcolor = img.getpixel((startx, starty))
      unfilled = colors[startcolor]
      nextfilled = set([startcoord])
   count = 1
   total = len(unfilled)
   while True:
      print("frame {}: {}/{} pixels".format(count, total - len(unfilled), total))
      frame = emptyframe.copy()
      oldfilled = nextfilled
      nextfilled = fillpixels(frame, unfilled, nextfilled, neighborfn, framestep)
      if len(oldfilled) > 0:
         xlist, ylist = zip(*list(oldfilled.union(nextfilled)))
         bbox = [
            max(0, min(xlist)),
            max(0, min(ylist)),
            min(max(xlist) + 1, maxx),
            min(max(ylist) + 1, maxy)
         ]
         print(bbox)
         yield (frame, bbox)
      else:
         return
      count += 1

# return a Graphics Control Extension per-frame header
def getgce():
   header = 0xf921 # technically two fields, 0x21 and 0xf9, in that order
   blocksize = 4
   disposal = 1 # don't dispose
   flag_userinput = 0
   flag_transparency = 1
   disposal_byte = disposal << 2 | flag_userinput << 1 | flag_transparency
   delay = 5 # 5ms between frames, most browsers won't go this low.
   transparent_idx = TRANSPARENT
   return struct.pack("<HBBHBB", header, blocksize, disposal_byte, delay, transparent_idx, 0)

def writeframes(img, frames, fp):
   gce = getgce()
   for s in GifImagePlugin.getheader(img)[0] + [gce] + GifImagePlugin.getdata(img):
      fp.write(s)
   for frame, bbox in frames:
      fp.write(gce)
      for s in GifImagePlugin.getdata(frame.crop(bbox), offset = bbox[:2]):
         fp.write(s)
   fp.write(";")

if __name__ == "__main__":
   parser = optparse.OptionParser()
   parser.add_option("-i", "--in", dest="infilename", default="img.gif",
      help="input filename (default img.gif)")
   parser.add_option("-o", "--out", dest="outfilename", default="out.gif",
      help="output filename (default out.gif)")
   parser.add_option("-s", "--start-coord", dest="startcoord", default=None,
      type="int", nargs=2,
      help="starting coordinate, in pixels from the upper-left-hand corner")
   parser.add_option("-d", "--diagonal", dest="diagonal", default=False,
      action="store_true", help="flood-fill diagonally")
   parser.add_option("-f", "--framestep", dest="framestep", default=1,
      type="int", help="number of fill steps per frame (default 1)")
   options, args = parser.parse_args()
   startcoord = None if options.startcoord is None else tuple(options.startcoord)
   neighborfn = moore_neighbors if options.diagonal else neighbors

   img = Image.open(options.infilename)
   img.load()
   setpalette(img)

   with open(options.outfilename, "wb") as fp:
      gifsicle = subprocess.Popen(["gifsicle", "-O3"], stdin=subprocess.PIPE, stdout=fp)
      frames = floodfill_frames(img, neighborfn, startcoord, options.framestep)
      writeframes(img, frames, gifsicle.stdin)
      print("Optimizing...")
      gifsicle.stdin.flush()
      gifsicle.stdin.close()
      gifsicle.wait()