3
\$\begingroup\$

I'd recommend thinking of readimage.py's tess() and cune() as sort of black boxes that use the OCRs. Anyway, this code is meant to be used for a science fair project where I am testing Tesseract and Cuneiforms' abilities to read text on images with various font sizes and colors, etc. Any thoughts, obvious mistakes, etc.?

main.py:

# Command-line arguments and other functionalities
import os
import sys
import math
import random
import ast
import argparse

# Image handling and OCR
import readimage
import drawimage
import distance

# Constants
DIMENSIONS = [850, 1100, 50, 50] # Width, Height, Side Margin, Top Margin
DICTLOC = "dict.txt"
COLORS = {
    "R" : ((255,0,0), "Red"),
    "G" : ((0,255,0), "Green"),
    "W" : ((255,255,255), "White"),
    "B" : ((0,0,0), "Black"),
    "Y1" : ((255,252,239), "Yellow1"),
    "Y2" : ((255,247,218), "Yellow2"),
    "Y3" : ((255,237,176), "Yellow3"),
    "Y4" : ((255,229,139), "Yellow4"),
}

# Read command-line arguments
parser = argparse.ArgumentParser()
parser.add_argument("-p", "--pages", type=int, help="Pages per Setting", default=1)
parser.add_argument("-f", "--fonts", help="Comma-Seperated List of fonts", default="freefont/FreeMono.ttf")
parser.add_argument("-tc", "--txtcolors", help="Comma-Seperated Color Initials", default="B")
parser.add_argument("-bc", "--bgcolors", help="Comma-Seperated Color Initials", default="W")
parser.add_argument("-hs", "--headsizes", type=str, help="Comma-Seperated Header Font Heights", default="50")
parser.add_argument("-bs", "--bodysizes", type=str, help="Comma-Serperated Body Font Heights", default="25")
parser.add_argument("-v", "--verbose", help="Print progress", action="store_true")
args = parser.parse_args()
pages = args.pages
fonts = args.fonts.split(",")
txtcolors = [COLORS[c] for c in args.txtcolors.split(",")]
bgcolors = [COLORS[c] for c in args.bgcolors.split(",")]
headsizes = [int(s) for s in args.headsizes.split(",")]
bodysizes = [int(s) for s in args.bodysizes.split(",")]
verbose = args.verbose

# Grab dictionary as list of words
worddict = open(DICTLOC).read()
worddict = worddict.split("\n")

def image_stats(file, correct, language="eng", tessconfig=""):
    tess = {}
    tess_out, tess["time"] = readimage.tess_ocr("img.png")
    tess_out = " ".join(tess_out.split()).strip()
    tess["dist"] = distance.lev(correct, tess_out)
    tess["per"] = round((len(correct)-tess["dist"])/len(correct),4)
    tess["tpc"] = round(tess["time"]/len(correct)*1000, 4)

    cune = {}
    cune_out, cune["time"] = readimage.cune_ocr("img.png")
    cune_out = " ".join(cune_out.split()).strip()
    cune["dist"] = distance.lev(correct, cune_out)
    cune["per"] = round((len(correct)-cune["dist"])/len(correct),4)
    cune["tpc"] = round(cune["time"]/len(correct)*1000, 4)
    return tess, cune

def main():
    if os.path.exists("fullout.txt"):
        os.remove("fullout.txt")
    if os.path.exists("avgout.txt"):
        os.remove("avgout.txt")
    fullout = open("fullout.txt",mode='a')
    fullout.write("\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\n")
    avgout = open("avgout.txt", mode='a')
    avgout.write("\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\n")
    for font in fonts:
        for txtcolor in txtcolors:
            for bgcolor in bgcolors:
                fullout.write(f"Font: {font}, {txtcolor[1]} on {bgcolor[1]}\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\n")
                avgout.write(f"Font: {font}, {txtcolor[1]} on {bgcolor[1]}\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\tCuneiform\tTesseract\n")
                for headsize in headsizes:
                    for bodysize in bodysizes:
                        cune_stats = []
                        tess_stats = []
                        avgout.write(f"{bodysize}")
                        for page in range(pages):
                            fullout.write(f"{bodysize}")
                            title = drawimage.generate_words(worddict, random.randint(1,10))
                            body = drawimage.generate_words(worddict, 10000)
                            img, correct = drawimage.create_page(title, body, DIMENSIONS, txtcolor[0], bgcolor[0], headsize, bodysize, font)
                            img.save("img.png")
                            correct = " ".join(correct).replace("\n", " ")
                            tess, cune = image_stats("img.png", correct)
                            tess_stats.append(tess)
                            cune_stats.append(cune)
                            fullout.write(f"\t{cune['time']}\t{tess['time']}\t{cune['tpc']}\t{tess['tpc']}\t{cune['dist']}\t{tess['dist']}\t{cune['per']}\t{tess['per']}\n")
                        cune = {}
                        tess = {}
                        for stat in cune_stats[0]:
                            cune[stat] = round(sum([i[stat] for i in cune_stats]) / len(cune_stats), 4)
                            tess[stat] = round(sum([i[stat] for i in tess_stats]) / len(tess_stats), 4)
                        avgout.write(f"\t{cune['time']}\t{tess['time']}\t{cune['tpc']}\t{tess['tpc']}\t{cune['dist']}\t{tess['dist']}\t{cune['per']}\t{tess['per']}\n")
    fullout.close()
    avgout.close()
if __name__ == "__main__":
    main()

drawimage.py:

from PIL import Image, ImageDraw, ImageFont
import random

# Turn words into lines, based on size of page and font, then return lines and height of lines
def word_space(words, font, height, spaceh=30):
    linew = 0
    linet = "" 
    lines = []
    wordnum = 0
    while wordnum < len(words):
        if len(linet) > 0: linet += " " 
        linet += words[wordnum]
        if font.getsize(linet)[0] > DIMENSIONS[0] - (2*DIMENSIONS[2]):
            if spaceh * (len(lines)+1) > height:
                linet = linet[:-(len(words[wordnum])+1)]
                break
            else:
                linet = linet[:-(len(words[wordnum])+1)]
                if font.getsize(words[wordnum])[0] > DIMENSIONS[0] - (2*DIMENSIONS[2]):
                    print("Word too long, skipping: " + words[wordnum])
                    wordnum += 1
                else:
                    lines.append(linet)
                    linet = ""
        else:
            wordnum += 1
    if linet:
        lines.append(linet)

    return lines, spaceh * len(lines)


# Add text to an image, return new image
def add_text(img, text, pos, font, fcolor):

    d = ImageDraw.Draw(img)
    d.text(pos, text, font=font, fill=fcolor)
    return img

# Draw an entire page, return image and correct text
def create_page(title, body, DIM, txtcolor, bgcolor, titlesize, bodysize, font):
    global DIMENSIONS
    DIMENSIONS = DIM
    img = Image.new('RGBA', (DIMENSIONS[0], DIMENSIONS[1]), bgcolor+(255,))
    titlefont = ImageFont.truetype(font, titlesize)
    bodyfont = ImageFont.truetype(font, bodysize)
    titlespaced, titleh = word_space(title, titlefont, DIMENSIONS[1]-40, spaceh=titlesize+10)
    for i, line in enumerate(titlespaced):
        img = add_text(img, line, (50,(titlesize+10)*i+20),titlefont,txtcolor)
    bodyspaced, margin = word_space(body, bodyfont, DIMENSIONS[1]-40-titleh-20, spaceh=bodysize+10)
    for i, line in enumerate(bodyspaced):
        img = add_text(img, line, (50,((bodysize+10)*i)+titleh+20), bodyfont, txtcolor)
    return img, titlespaced+bodyspaced

# Generate and return a given number of words
def generate_words(worddict, length):
    words = []
    for j in range(length):
        word = random.choice(worddict)
        mod = random.randint(1,10)
        if mod == 1:
            word = word.upper()
        elif mod == 2:
            word = word.capitalize()
        elif random.randint(1,15) == 1:
            word += "."
        words.append(word)
    return words

readimage.py:

import subprocess
import os
import time
from PIL import Image
import pytesseract

# Functions to run either OCR on a given image.

def tess_ocr(file, language="eng", config=""):
    # Run and time Tesseract, return output
    start = time.time()
    out = pytesseract.image_to_string(Image.open(file), language, config=config)
    return out, round(time.time() - start, 4)

def cune_ocr(file, language="eng"):
    # Run Cuneiform on image
    start = time.time()
    subprocess.call(["cuneiform", "-o", "cuneout.txt",file], stdout=subprocess.PIPE)

    # Fetch and return output
    if os.path.exists("cuneout.txt"):
        out = open("cuneout.txt").read()
        return out, round(time.time() - start, 4)
    else:
        print("Cuneiform reported no output, returning empty string")
        return "", round(time.time() - start, 4)

distance.py:

import numpy

# Find Levenshtein Distance between two strings
def lev(a,b):
    sizex = len(a)+1
    sizey = len(b)+1
    matrix = numpy.zeros((sizex,sizey))
    for x in range(sizex):
        matrix[x,0] = x
    for y in range(sizey):
        matrix[0,y] = y

    for y in range(1, sizey):
        for x in range(1, sizex):
            cost = 0
            if a[x-1] != b[y-1]:
                cost = 2
            matrix[x,y] = min(
                matrix[x-1,y] + 1,
                matrix[x,y-1] + 1,
                matrix[x-1,y-1] + cost
            )

    return int(matrix[sizex-1,sizey-1])
\$\endgroup\$
2
\$\begingroup\$

Some small comments:

  • In main.py: argparse can already deal with argument types and multiple arguments (it even enforces the type across multiple arguments). If you change it to multiple arguments being separated by whitespace, you can simply use this:

    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--pages", type=int, help="Pages per Setting", default=1)
    parser.add_argument("-f", "--fonts", nargs="+", help="Space seperated List of fonts", default=["freefont/FreeMono.ttf"])
    parser.add_argument("-tc", "--txtcolors", nargs="+", help="Space seperated Color Initials", default=["B"])
    parser.add_argument("-bc", "--bgcolors", nargs="+", help="Space seperated Color Initials", default=["W"])
    parser.add_argument("-hs", "--headsizes", nargs="+", type=int, help="Space seperated Header Font Heights", default=[50])
    parser.add_argument("-bs", "--bodysizes", nargs="+", type=int, help="Space seperated Body Font Heights", default=[25])
    parser.add_argument("-v", "--verbose", help="Print progress", action="store_true")
    args = parser.parse_args()
    
  • In readimage.py: If you need something more than once, don't repeat yourself. Write a wrapper function that does the timing for you:

    from time import perf_counter
    from functools import wraps
    
    def timeit(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = perf_counter()
            ret = f(*args, **kwargs):
            return ret, round(perf_counter() - start, 4)
        return wrapper
    
    
    @timeit
    def tess_ocr(file, language="eng", config=""):
        return pytesseract.image_to_string(Image.open(file), language, config=config)
    
    @timeit
    def cune_ocr(file, language="eng"):
        # Run Cuneiform on image
        subprocess.call(["cuneiform", "-o", "cuneout.txt",file], stdout=subprocess.PIPE)
    
        # Fetch and return output
        if os.path.exists("cuneout.txt"):
            with open("cuneout.txt") as f:
                return f.read()
        else:
            print("Cuneiform reported no output, returning empty string")
            return ""
    

    Note that I used the more accurate time.perf_counter and ensured that the docstring etc are conserved using functools.wraps. I also made sure the file is properly closed using a with statement.

    Arguably, the wrapper function should only print/log the timing instead of returning it and thereby changing the signature of the function, but I left your interface as is here. You might even want to put this function into its own module, since it is quite reusable (especially when modifying it not to change the function signature).

  • There are already Levenshtein distance modules, like editdistance and python-Levenshtein, both of which are quite fast.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.