ocrotest-models

#!/usr/bin/python

import sys,os,re,glob,math,glob,signal,cPickle
import iulib,ocropus
import ocropy
from pylab import *
from optparse import OptionParser
import ocrolib
from ocrolib.fgen import *

parser = OptionParser(usage="""
%prog [model n]

This is a "smoke test" for a set of classifiers.  It generates a few dozen images of 30 different
characters (including umlaut, Greek, and Japanese) and then trains and tests each character
classifier on those.  The tests should complete without crashing, and they should usually have
zero errors ("OK"), although a small number of errors may be OK.

If exactly two arguments are given, the first is used to construct the classifier (it may be
an OCRopus component or a Python construction expression), and the second one is the number
of training examples per class.
""")
parser.add_option("-D","--display",help="display continuously",action="store_true")
parser.add_option("-n","--nvariants",help="number of variants",type=int,default=20)
options,args = parser.parse_args()

if len(args)==0:
    args = ["=ocrolib.AutoMlpModel()"]

models = []
for arg in args:
    model = ocrolib.load_component(arg)
    models.append((model,options.nvariants))

class unif:
    def __init__(self,lo,hi):
        self.lo = lo
        self.hi = hi
    def __call__(self):
        return rand()*(self.hi-self.lo)+self.lo

scale = unif(2.0,2.0)
aspect = unif(0.95,1.05)
margin = unif(0.0,1.0)
noise = unif(0.0,0.1)
maxdelta = unif(0.0,1.0)
rotation = unif(-1.0,1.0) # in degrees
inner = unif(1.0,1.0)
testset = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")+[u"\u03c8",u"\u00d6","FF",u"\u30DD"]
spec = "Arial"
size = 30
nvariants = 20
display = options.display

def generate(nvariants):
    for d in range(nvariants):
        for i in range(len(testset)):
            cls = testset[i]
            image = pango_render_string(cls,spec=spec,size=size,pad=10,markup=0,
                                        scale=scale(),aspect=aspect(),rotation=rotation())
            image = average(array(image,'f'),axis=2)
            image /= amax(image)
            image = gauss_degrade(image,margin=margin(),noise=noise(),inner=inner())
            mdl = maxdelta()
            image = gauss_distort([image],maxdelta=mdl)[0]
            yield image,cls

# quick test of the character encoding in IModel
# FIXME: move this to global functions
model = ocropy.make_IModel("AutoMlpClassifier")
assert model.chr(model.ord("A"))=="A"
assert model.chr(model.ord("AA"))=="AA"
assert model.chr(model.ord(u"\u00d6"))==u"\u00d6"
assert model.chr(model.ord(u"\u03c8"))==u"\u03c8"
assert model.chr(model.ord(u"\u30dd"))==u"\u30dd"

for model,nvariants in models:
    print "===",model.name(),nvariants,"==="
    model.setExtractor("ScaledImageExtractor")
    if display: ion(); show()
    print "generating"
    count = 0
    for image,cls in generate(nvariants):
        if display and count<=36: subplot(6,6,count%36+1); imshow(image); draw()
        model.cadd(image,cls)
        count += 1
    print "training"
    model.updateModel(verbose=1)
    print "saving"
    ocrolib.save_component("_temp.model",model)
    print "loading"
    model = ocrolib.load_component("_temp.model")
    print "testing"
    errs = 0
    for image,cls in generate(1):
        result = sorted(model.coutputs(image),key=lambda x:x[1],reverse=1)
        pred = result[0][0]
        if cls!=pred:
            errs += 1
            print "ERR",cls,result[:3]
    if errs>0: print "ERRS",errs
    else: print "OK"