pdfxml2csv

#!/usr/bin/python
# *-* coding: utf-8 *-*
import csv
import argparse
import operator
import numpy
import math
import lxml.etree
import lxml.html
from pprint import pprint as pprint_orig


def pprint(obj, *args, **kwargs):
    width = kwargs.pop('width', 120)
    pprint_orig(obj, *args, width=width, **kwargs)


def parse_xml(filename):
    pdfxml = open(filename, 'r').read()
    root = lxml.etree.fromstring(pdfxml)

    pages = []
    for pagenum, page in enumerate(root):
        assert page.tag == 'page'
        pageinfo = {
            'pagenumber': page.attrib.get('number'),
            'width': page.attrib.get('width'),
            'height': page.attrib.get('height'),
        }
        textboxes = []
        for v in page:
            if v.tag == 'text':
                left = int(v.attrib.get('left'))
                top = int(v.attrib.get('top'))
                width = int(v.attrib.get('width'))
                height = int(v.attrib.get('height'))
                info = {
                    'left': left,
                    'top': top,
                    'width': width,
                    'height': height,
                    'text': v.text,
                    }
                textboxes.append(info)
        pages.append((pageinfo, textboxes))
    return pages


def pages2lines(pages):
    rows = []
    # first create lines from given XML
    for pagenum, page in enumerate(pages):
        pagelines = {}
        for item in page[1]:
            top = item['top']
            # fix some off-by-one placement issues, which make some text span over two lines where it should be in one
            if (top - 1) in pagelines:
                top = top - 1
            elif (top + 1) in pagelines:
                top = top + 1
            line = pagelines.setdefault(top, [])
            line.append(item)

        ordered = list(sorted([(k, sorted(v, key=operator.itemgetter('left'))) for k, v in pagelines.iteritems()]))
        rows.extend(ordered)

    return rows


def lines2cols(rows, splits, header=-1, footer=-1, verbose=False):
    records = []
    # splits must containt catchall zero
    if 0 not in [i[0] for i in splits]:
        splits = [(0, 'left')] + splits
    splits = list(sorted(splits))

    for topoffset, line in rows:
        if footer != -1 and topoffset > footer:
            continue
        if header != -1 and topoffset < header:
            continue
        linerec = [[] for i in xrange(len(splits))]

        for item in sorted(line, reverse=True):
            left = item['left']

            idx = len(splits) - 1
            while splits[idx][0] > left:
                idx = idx - 1
            linerec[idx].insert(0, (left, item))
        if verbose:
            pprint(linerec)

        ordered = [sorted(i) for i in linerec]
        records.append(ordered)
    return records


def records2csv(table, csv_filename, csv_sep, meta):
    csvfd = open(csv_filename, 'wb')
    csvw = csv.writer(csvfd)
    for row in table:
        row2 = []
        for idx, key in meta:
            try:
                fieldinfo = row[idx][0][1]
                row2.append(str(fieldinfo[key]))
            except IndexError:
                row2.append('')

        for r in row:
            row2.append(csv_sep.join([i[1]['text'] for i in r]))
        csvw.writerow([i.encode('utf-8') for i in row2])
    csvfd.close()


def find_gaps(array, threshold=0):
    """finds "gaps", that is runs of values of less or equal to threshold in 1D numpy array"""
    gap_start = 0
    prev = None
    n = 0
    gaps = []
    while n < len(array):
        cur = array[n]
        if prev is None:
            prev = cur
            continue
        if prev > threshold and cur <= threshold:
            gap_start = n

        if prev <= threshold and cur > threshold:
            # end gap
            gaps.append((gap_start, n - gap_start))
        prev = cur
        n += 1

    if cur <= threshold and prev <= threshold:
        gaps.append((gap_start, n - gap_start))
    return gaps


def estimate_cell_vertical(pages, wlen=6):
    pagegaps = []
    for pageinfo, page in pages:
        # pagepoints = [0]*(int(float(pageinfo['height']))+1)
        pagepoints = numpy.zeros(int(float(pageinfo['height'])) + 1)
        for item in page:
            for idx in xrange(item['top'], item['top'] + item['height']):
                pagepoints[idx] += len(item['text'])
        window = numpy.hamming(wlen)
        smoothed = numpy.convolve(pagepoints, window)
        gaps = find_gaps(smoothed)
        pagegaps.append(gaps)
    return pagegaps


def estimate_columns(pages, wlen=6):

    maxwidth = max([int(math.ceil(float(pageinfo['width']))) for pageinfo, page in pages])
    counts = numpy.zeros(maxwidth, numpy.uint32)

    for pageinfo, page in pages:
        for item in page:
            if item['text'].strip():
                counts[item['left']:item['left'] + item['width']] += 1

    window = numpy.hamming(wlen)
    smoothed = numpy.convolve(counts, window) / 3
    avg = sum(counts) / len(counts)
    gaps = find_gaps(smoothed, threshold=avg * 0.1)

    pos = [(i[0] + i[0] + i[1]) / 2 for i in gaps if i[0] != 0]

    print 'Recommended --left parameter:'
    print '--left %s' % ','.join([str(i) for i in pos])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Parse PDF XML into a CSV')
    parser.add_argument('filename', metavar='filename.xml', type=str, nargs=1, help='input XML')
    parser.add_argument('--verbose', dest='verbose', action='store_true', default=False, help='verbose output')
    parser.add_argument('--header', dest='header', type=int, nargs='?', default=-1, help='exclude all content on each page above this')
    parser.add_argument('--footer', dest='footer', type=int, nargs='?', default=-1, help='exclude all content on each page below this')
    parser.add_argument('--left', dest='left', type=str, nargs='?', help='position of a new cell split')
    parser.add_argument('--meta', dest='meta', type=str, nargs='?', help='include metadata about fields')
    parser.add_argument('--estimate-vcell', dest='estimate_vcell', action='store_true', help='perform a vertical cell estimation')
    parser.add_argument('--estimate-columns', dest='estimate_columns', action='store_true', help='perform a column separator estimation')
    parser.add_argument('--csv', dest='output', type=str, nargs='?', help='output CSV file name')
    parser.add_argument('--csv-textbox-separator', dest='csv_sep', type=str, default=u' ', nargs='?', help='output CSV separator for joining PDF TextBox elements')

    args = parser.parse_args()
    verbose = args.verbose
    if args.left:
        leftsplits = list(sorted([(int(i.strip()), 'left') for i in args.left.split(',')]))
    else:
        leftsplits = [(0, 'left')]

    pages = parse_xml(args.filename[0])
    rows = pages2lines(pages)
    table = lines2cols(rows,
        header=args.header,
        footer=args.footer,
        splits=leftsplits,
        verbose=verbose)

    if args.estimate_columns:
        gaps = estimate_columns(pages)

    if args.estimate_vcell:
        gaps = estimate_cell_vertical(pages)

    meta = []
    if args.meta:
        for item in args.meta.split(','):
            if '-' not in item:
                raise ValueError("invalid meta var name")
            field, key = item.split('-')
            if key not in ('left', 'width', 'top', 'height'):
                raise ValueError("invalid meta var name")
            try:
                field = int(field)
            except ValueError:
                raise ValueError("invalid meta field ref index")
            meta.append((field, key))

    if args.output:
        records2csv(table, args.output, args.csv_sep, meta)