#!/usr/bin/env python
"""
extracts the original (sharp) image and depth map from images shot with
Google's Android camera application in "Lens Blur" mode
"""
import sys, os, re

if __name__ == "__main__":
    try:
        infile = sys.argv[1]
    except IndexError:
        print >>sys.stderr, "Usage:", sys.argv[0], "<input.jpg>"
        sys.exit(2)

    def errexit(msg):
        print >>sys.stderr, "ERROR:", msg
        sys.exit(1)

    # extract the XMP metadata from the file by walking the APP1 marker chain
    try:
        f = open(infile, "rb")
    except IOError, e:
        errexit("could not open input file - %s" % e)
    xmp = ""
    try:
        if f.read(2) != "\xff\xd8":
            errexit("input file is not a JPEG file")
        while True:
            marker = f.read(2)
            if not marker.startswith("\xff"):
                errexit("marker expected, but not found")
            if marker != "\xff\xe1":
                break
            length = f.read(2)
            if len(length) != 2:
                errexit("unexpected EOF while reading marker length")
            length = ord(length[0]) * 256 + ord(length[1]) - 2
            if length < 0:
                errexit("invalid marker length")
            data = f.read(length)
            if len(data) != length:
                errexit("unexpected EOF while reading marker")
            try:
                marker, data = data.split('\0', 1)
            except ValueError:
                errexit("invalid APP1 marker")
            if marker.startswith("http://ns.adobe.com/xap/"):
                xmp = data
            elif marker.startswith("http://ns.adobe.com/xmp/extension/"):
                xmp += data[40:]   # not checking the header
            else:
                pass  # Exif -- not relevant for us
    finally:
        f.close()

    # search for the images in the XMP file
    for itype, idata in re.findall(r'g(depth|image):data="([^"]*)"', xmp, re.S + re.I):
        try:
            idata = idata.decode('base64')
        except Exception:
            print >>sys.stderr, "ERROR:", itype, "is not valid base64"
            continue
        if idata.startswith("\xff\xd8"):
            ext = "jpg"
        elif idata.startswith("\x89PNG"):
            ext = "png"
        else:
            ext = "bin"
        outfile = "%s_%s.%s" % (os.path.splitext(infile)[0], itype, ext)
        try:
            f = open(outfile, "wb")
            f.write(idata)
            f.close()
        except IOError, e:
            print >>sys.stderr, "ERROR: could not write output file '%s' -" % outfile, e