#!/usr/bin/env python
# $URL: https://pypng.googlecode.com/svn/trunk/code/pipdither $
# $Rev: 100 $

# pipdither
# Error Diffusing image dithering.
# Now with serpentine scanning.

# See http://www.efg2.com/Lab/Library/ImageProcessing/DHALF.TXT

import png

def dither(out, inp, bitdepth=1, linear=False, defaultgamma=None):
    """Dither the input PNG `inp` into an image with a smaller bit depth
    and write the result image onto `out`.  `bitdepth` specifies the bit
    depth of the new image.
    
    Normally the source image gamma is honoured (the image is
    converted into a linear light space before being dithered), but
    if the `linear` argument is true then the image is treated as
    being linear already: no gamma conversion is done (this is
    quicker, and if you don't care much about accuracy, it won't
    matter much).  Normally images with no gamma indication (no
    ``gAMA`` chunk) are treated as linear (gamma = 1.0), but often
    it can be better to assume a different gamma value: For example
    continuous tone photographs intended for presentation on the
    web often carry an implicity assumption of being encoded with
    a gamma of about 0.45 (because that's what you get if you just
    "blat the pixels" onto a PC framebuffer), so ``defaultgamma=0.45``
    might be a good idea.  `defaultgamma` does not override a gamma
    value specified in the file itself: It is only used when the
    file does not specify a gamma.
    
    If you (pointlessly) specify both `linear` and `defaultgamma`,
    `linear` wins.
    
    """

    # The dithering algorithm is not completely general; it
    # can only do bit depth reduction, not arbitrary palette changes.
    import operator
    maxval = float(2**bitdepth - 1)
    r = png.Reader(file=inp)
    # If image gamma is 1 or gamma is not present and we are assuming a
    # value of 1, then it is faster to pass a maxval parameter to
    # asFloat (the multiplications get combined).  With gamma, we have
    # to have the pixel values from 0.0 to 1.0 (as long as we are doing
    # gamma correction here).
    # Slightly annoyingly, we don't know the image gamma until we've
    # called asFloat().
    _,_,pixels,info = r.asFloat()
    planes = info['planes']
    assert planes == 1
    width = info['size'][0]
    if linear:
        gamma = None
    else:
        gamma = info.get('gamma') or defaultgamma
    if gamma == 1:
        # This is equivalent to the no conversion case.
        gamma = None
    # Convert gamma from encoding gamma to the required power to convert
    # from image pixel values to a linear intensity space.
    if gamma:
        gamma = 1.0/gamma
    def iterdither():
        ed = [0.0]*width
        flipped = False
        for row in pixels:
            if gamma:
                row = map(gamma.__rpow__, row)
            row = map(maxval.__mul__, row)
            row = map(operator.add, ed, row)
            if flipped:
                row = row[::-1]
            targetrow = [0] * width
            for i,v in enumerate(row):
                # Clamp.  Necessary because previously added errors may take
                # v out of range.
                v = max(0, min(v, maxval))
                # t is the chosen target value, as an int.
                t = int(round(v))
                targetrow[i] = t
                # err is the error that needs distributing.
                err = v - t
                # Sierra "Filter Lite" distributes          * 2
                # as per this diagram.                    1 1
                ef = err/2.0
                # :todo: consider making rows one wider at each end and
                # removing "if"s
                if i+1 < width:
                    row[i+1] += ef
                ef /= 2.0
                ed[i] = ef
                if i:
                    ed[i-1] += ef
            if flipped:
                ed = ed[::-1]
                targetrow = targetrow[::-1]
            yield targetrow
            flipped = not flipped
    info['bitdepth'] = bitdepth
    w = png.Writer(**info)
    w.write(out, iterdither())


def main(argv=None):
    # http://www.python.org/doc/2.4.4/lib/module-getopt.html
    from getopt import getopt
    import sys
    if argv is None:
        argv = sys.argv
    opt,argv = getopt(argv[1:], 'a:b:l')
    k = {}
    for o,v in opt:
        if o == '-a':
            k['defaultgamma'] = float(v)
        if o == '-b':
            k['bitdepth'] = int(v)
        if o == '-l':
            k['linear'] = True
        if o == '-?':
            print >>sys.stderr, "pipdither [-a assumed-gamma] [-b bits] [-l] [in.png]"

    if len(argv) > 0:
        f = open(argv[0], 'rb')
    else:
        f = sys.stdin

    return dither(sys.stdout, f, **k)


if __name__ == '__main__':
    main()
