ocroskew-pred

#!/usr/bin/python
import os
import os.path
import argparse

import ocrorot
import scipy.ndimage as ndi
from pylab import *
from torch import nn
from dlinputs import gopen, paths, utils, filters

model_path = os.environ.get(
    "MODELS", ".:/usr/local/share/ocrorot:/usr/share/ocrorot")
default_model = "logskew-000015808-000132.pt"

parser = argparse.ArgumentParser("train a page segmenter")
parser.add_argument("-m", "--model", default=default_model, help="load model")
parser.add_argument("-b", "--batchsize", type=int, default=1)
parser.add_argument("-D", "--makesource", default=None)
parser.add_argument("-P", "--makepipeline", default=None)
parser.add_argument("-i", "--invert", action="store_true")
parser.add_argument("--display", type=int, default=0)
parser.add_argument("input")
parser.add_argument("output", nargs="?")

args = parser.parse_args()
ARGS = {k: v for k, v in args.__dict__.items()}

if args.display > 0:
    rc("image", cmap="gray")
    ion()


def make_source():
    return gopen.sharditerator_once(args.input)


def make_pipeline():

    def fixdepth(image):
        assert image.ndim in [2, 3]
        if image.ndim == 3:
            image = np.mean(image, 2)
        image -= amin(image)
        image /= amax(image)
        if args.invert:
            image = 1-image
        return image

    return filters.compose(
        filters.rename(input="bin.png png"),
        filters.map(input=fixdepth))


if args.makesource:
    execfile(args.makesource)
if args.makepipeline:
    execfile(args.makepipeline)


def pixels_to_batch(x):
    b, d, h, w = x.size()
    return x.permute(0, 2, 3, 1).contiguous().view(b*h*w, d)


class PixelsToBatch(nn.Module):
    def forward(self, x):
        return pixels_to_batch(x)


source = make_source()
pipeline = make_pipeline()
source = pipeline(source)
if args.output:
    sink = gopen.open_sink(args.output)

mname = paths.find_file(model_path, args.model)
assert mname is not None, "model not found"
print "loading", mname

sk = ocrorot.SkewEstimator(mname)
print sk.model


def display_batch(image, output):
    clf()
    if image is not None:
        subplot(121)
        imshow(image[0, :, :, 0], vmin=0, vmax=1)
    if output is not None:
        subplot(122)
        imshow(output[0, :, :, 0], vmin=0, vmax=1)
    draw()
    ginput(1, 1e-3)


for i, sample in enumerate(source):
    fname = sample["__key__"]
    image = sample["input"]
    output = sk.skew(image)
    print i, fname, output
    if args.output:
        if angle != 0.0:
            corrected = ndi.rotate(image, -output, order=1)
        else:
            corrected = image
        result = utils.metadict(sample, {
            "__key__": fname,
            "bin.png": corrected
        })
        sink.write(result)

if args.output:
    sink.close()