fpicencode.py

#!/usr/bin/python

'''Encodes a file into a paletted image.

NAME
  fpicencode

SYNOPSIS
  fpicencode [-z] [-m sourceimage] sourcefile [outputfile]

DESCRIPTION
  The fpicencode script encodes an arbtrary file as a roughly square
  256-colour image.

  Options:

  -z        Store the file compressed, if the image size will be reduced
            as a result.

  -m image  Use an image to be modulated by the data at the cost of
            halving the data density.

  This script can be found on Pastebin.
  Decoder: http://pastebin.com/YKSSX9n7
  Encoder: http://pastebin.com/mysKcugD

'''

import sys
import os.path
from math import *
from PIL import Image
import zlib

gamma = 2.2

class Error (Exception):
  pass

def vlengthsquared(a):
  return sum(x * x for x in a)

def vlength(a):
  return sqrt(vlengthsquared(a))

def vneg(a):
  return tuple(-x for x in a)

def vsum(*vectorargs):
  if len(vectorargs) == 1:
    vectors = vectorargs[0]
  else:
    vectors = vectorargs
  r = tuple(vectors[0])
  for i in range(1, len(vectors)):
    r = tuple(a + b for a, b in zip(r, vectors[i]))
  return r

def vdiff(a, b):
  return tuple(x - y for x, y in zip(a, b))

def vscaled(a, scale):
  return tuple(x * scale for x in a)

def vlerp(a, b, t):
  return vsum(
    vscaled(a, 1.0 - t),
    vscaled(b, t)
  )

def fabulouspalette():
  anchors = [
    # Cube (first half)
    (0, 0, 0), (0, 0, 254), (254, 0, 0), (254, 0, 254),
    # Octahedron
    (127, 127, 0), (127, 0, 127), (0, 127, 127),
    (127, 127, 254), (127, 254, 127), (254, 127, 127),
    # Centre
    (127, 127, 127),
    # Dark grey
    (63, 63, 63),
    # Cube (last half)
    (0, 254, 0), (0, 254, 254), (254, 254, 0), (254, 254, 254)
  ]
  anomalousexcursions = {
    (-1, -1, -1): (1, 1, 2),
    (-1, -1, 256): (1, 1, 253),
    (256, -1, -1): (254, 1, 2),
    (256, -1, 256): (254, 1, 253),
    (-1, 256, -1): (1, 254, 2),
    (-1, 256, 256): (1, 254, 253),
    (256, 256, -1): (254, 254, 2),
    (256, 256, 256): (254, 254, 253),
  }
  p = [0] * (3 * 256)
  for maincolix in range(16):
    a = anchors[maincolix]
    for lix in range(2):
      lo = -lix
      ls = 1 + 2 * lix
      for gix in range(2):
        for rix in range(2):
          for bix in range(2):
            am = (rix, gix, bix)
            ucol = tuple(a[i] + lo + ls * am[i] for i in range(3))
            if ucol in anomalousexcursions:
              ucol = anomalousexcursions[ucol]
            col = [max(0, min(255, uc)) for uc in ucol]
            subcolix = (gix << 3) + (rix << 2) + (bix << 1) + lix
            pix = 3 * ((maincolix << 4) + subcolix)
            p[pix : pix + 3] = col
  return p

def gammac(level):
  return pow(level / 255.0, gamma)

def rup2(x):
  return (x + 1) & ~1

def chunksizebytes(datasize):
  r = bytearray(4)
  for i in range(4):
    r[3 - i] = (datasize >> (8 * i)) & 255
  return r

def chunkidbytes(chunkid):
  r = bytearray(4)
  for i in range(4):
    r[i] = ord(chunkid[i])
  return r

def chunkhdr(chunkid, datasize):
  return chunkidbytes(chunkid) + chunksizebytes(datasize)

def newchunk(chunkid, data):
  r = chunkhdr(chunkid, len(data)) + data
  if len(data) & 1:
    a = bytearray(1)
    a[0] = 0
    r += a
  return r

def chunksize(chunkhdr):
  r = 0
  for i in range(4):
    r += (chunkhdr[7 - i] & 0x00FF) << (8 * i)
  r = 4 + rup2(r)
  return r

def main():

  infname = None
  picfname = None
  outfname = 'encoded.png'
  dozip = False
  iszipped = False
  domodulate = False
  fname = infname
  rc = 0
  args = list(sys.argv)
  cn = args[0]
  opts = []
  params = []
  argix = 1
  while argix < len(args):
    arg = args[argix]
    nextarg = args[argix + 1] if argix + 1 < len(args) else None
    if len(params) == 0 and len(arg) >= 2 and arg[0] == '-':
      opt = arg[:2]
      hassep = len(arg) >= 3 and arg[2] in ':='
      if opt == '-z':
        opts.append((opt, None))
      elif opt == '-m':
        if hassep:
          opts.append((opt, arg[3:]))
        else:
          opts.append((opt, nextarg))
          argix += 1
      else:
        opts.append(arg)
    else:
      params.append(arg)
    argix += 1
  for opt, value in opts:
    if opt == '-z':
      dozip = True
    elif opt == '-m':
      picfname = value
      domodulate = picfname is not None and picfname != ''
  if len(params) >= 1:
    infname = params[0]
    params = params[1:]
  if len(params) >= 1:
    outfname = params[0]
    params = params[1:]

  if infname is None or infname == '':
    print cn + ': A source filename is expected.'
    rc = 1
  elif len(params) > 0:
    print cn + ': Too many parameters.'
    rc = 1

  if rc:
    return rc

  fname = os.path.split(infname)[1]
  fdata = None
  fh = open(infname, 'rb')
  try:
    rc = 2
    fdata = bytearray(fh.read())
    rc = 0
  finally:
    fh.close()
  if rc:
    return rc

  cdata = fdata

  if dozip:
    zdata = zlib.compress(str(cdata), 9)
    if len(zdata) < len(cdata):
      cdata = bytearray(zdata)
      iszipped = True

  hdr = bytearray([0] * 20)
  hdr[0] = 1 if iszipped else 0

  ct = chunkidbytes('Raw ')
  hc = newchunk('RDHD', hdr)
  nc = newchunk('NAME', bytearray(fname + chr(0)))
  dc = newchunk('DATA', cdata)
  iffdata = ct + hc + nc + dc
  iffdatasize = len(iffdata)

  edata = chunkidbytes('FORM') + chunksizebytes(iffdatasize) +  iffdata;

  #fh = open('data.iff', 'wb')
  #try:
  #  fh.write(str(edata))
  #finally:
  #  fh.close()

  mdata = edata

  if domodulate:
    mdata = bytearray(2 * len(edata))
    for i in range(len(edata)):
      e = edata[i]
      mdata[2 * i] = (e >> 4) & 15
      mdata[2 * i + 1] = e & 15

  imgfmt = 'P'
  palette = fabulouspalette()

  reqpix = len(mdata)
  w = 8 * int(round(sqrt(reqpix * 4.0/3.0) / 8.0))
  h = w * 3 // 4
  while w * h < reqpix:
    w += 8
    h = w * 3 // 4
  im = Image.new(imgfmt, (w, h))
  im.putpalette(palette)

  pic = None
  pixrem = w * h;
  x = 0
  y = 0
  mix = 0

  if domodulate:
    pic = Image.open(picfname).convert('RGB').resize((w, h), Image.CUBIC)
    gcpalette = tuple(gammac(c) for c in palette)
    errshares = [rw / 16.0 for rw in [7, 3, 5, 1]]
    fwderr = None
    rowerrs = None
    nextrowerrs = [(0.0, 0.0, 0.0)] * (w + 2)
    while pixrem:
      if x == 0:
        fwderr = (0.0, 0.0, 0.0)
        rowerrs = nextrowerrs
        nextrowerrs = [(0.0, 0.0, 0.0)] * (w + 2)
      m = mdata[mix]
      nlsrcpx = pic.getpixel((x, y))
      srcpx = tuple(gammac(c) for c in nlsrcpx)
      idealcol = vdiff(srcpx, vsum(fwderr, rowerrs[1 + x]))
      bestcix = i << 4
      beste2 = 1e9
      for i in range(16):
        cix = (i << 4) + (m & 15)
        col = gcpalette[3 * cix : 3 * cix + 3]
        e2 = vlengthsquared(vdiff(col, idealcol))
        if e2 < beste2:
          bestcix = cix
          beste2 = e2
      cix = bestcix
      col = gcpalette[3 * cix : 3 * cix + 3]
      im.putpixel((x, y), cix)
      err = vdiff(col, idealcol)
      fwderr = vscaled(err, errshares[0])
      nextrowerrs[x] = vsum(nextrowerrs[x], vscaled(err, errshares[1]))
      nextrowerrs[x + 1] = vsum(nextrowerrs[x + 1], vscaled(err, errshares[2]))
      nextrowerrs[x + 2] = vsum(nextrowerrs[x + 2], vscaled(err, errshares[3]))
      mix = (mix + 1) % len(mdata)
      x = (x + 1) % w
      y += 0 if x else 1
      pixrem -= 1
  else:
    while pixrem:
      px = mdata[mix]
      im.putpixel((x, y), px)
      mix = (mix + 1) % len(mdata)
      x = (x + 1) % w
      y += 0 if x else 1
      pixrem -= 1

  ext = os.path.splitext(outfname)[1]
  fmt = 'GIF' if ext in ['.gif', '.GIF'] else 'PNG'

  im.save(outfname, fmt)

  return rc

if __name__ == '__main__':
  sys.exit(main())