dump_reader.py

#!/usr/bin/env python3

import cv2 as cv
# from matplotlib import pyplot as plt
import math
import numpy as np
import tesserocr
from PIL import Image
import statistics
import sys

tes = tesserocr.PyTessBaseAPI(path='/opt/homebrew/share/tessdata/',
                              lang='ohfm', # change for finetuning
                              psm=tesserocr.PSM.SINGLE_LINE)
res = tes.SetVariable('tessedit_char_whitelist', '0123456789ABCDEFSUM: ')
assert res == True

#imagenum = 0

"""
targets = [None]*16
bf = cv.BFMatcher()
detector = cv.FeatureDetector_create('SIFT')
extractor = cv.DescriptorExtractor_create('SIFT')
for i in range(len((targets))):
    im = cv.imread(f'train_{i:x}.png', cv.IMREAD_GRAYSCALE)
    targets[i] = detector.detect(im, None)

def identify_char(im):
    global targets
    global d = etector
    dists = [None]*16
    low = float('inf')
    idx = None
    for i, t in enumerate(targets):
        _, des = detector.detectAndCompute(im, None)
        matches =  bf.knnMatch(t, des, k=2)
        dists = [m.distance for m in matches]
        good = []
        for m, n in matches:
            if m.distance < 0.75*n.distance:
                good.append([m])

        d = sum(dists)/len(dists)
        if d < low:
            low = d
            idx = i
    return f'{idx:X}'
"""

def filter_average(avg, t, wnoise):
    # start with gaps, then remove those to yield the segments of text
    gaps = find_gaps(avg, t, wnoise)
    start = 0
    l = []
    for g in gaps:
        end = g[0]
        # don't include empty spans, which could show up at the endpoints
        if (end - start) > wnoise:
            l.append((start, end))
        start = g[1]
    end = len(avg)
    # this may have left us with boxes containing multiple characters.
    # to address this, we can test the width of each box and subdivide if
    # it's too wide
    med = statistics.median(x[1] - x[0] for x in l)
    # print(f"median: {med}")
    out = []
    for r in l:
        width = r[1] - r[0]
        # print(f"width: {width}")
        splits = int(width/med)
        if splits <= 1:
            # print(f'adding {r}')
            out.append(r)
        else:
            chunk = width/splits
            # print(f"chunk: {chunk}")
            for i in range(splits):
                start = int(round(r[0] + i*chunk))
                end = int(round(r[0] + (i+1)*chunk))
                # print(f'adding split {(start, end)}')
                out.append((start, end))
    return out

def find_gaps(avg, t, wnoise):
    ranges = []
    i = 0
    start = 0
    end = 0
    i = 0
    while i < len(avg):
        # continue the current span until we hit a dark point
        if avg[i] < t:
            end = i
            if (end - start) != 0:
                ranges.append((start, end))
            # skip past the dark section
            while i < len(avg) and avg[i] < t:
                i += 1
            start = i
        else:
            i += 1
    # end the last span if it wasn't yet ended
    if i != end:
        ranges.append((start, end))
    # merge ranges around small interruptions
    if len(ranges) < 2:
        return ranges
    out = []
    i = 0
    while i < len(ranges):
        start, end = ranges[i]
        i += 1
        # update the end of the old range to include next range if the gap
        # between them is small enough
        while i < len(ranges) and (ranges[i][0] - end) < wnoise:
            end = ranges[i][1]
            i += 1
        out.append((start, end))
    return out

def process_box(img):
    # produce row vector averages
    row_mean = np.mean(img, axis=1)
    rows = filter_average(row_mean, 240, 20)
    # print('column math')
    col_mean = np.mean(img, axis=0)
    cols = filter_average(col_mean, 240, 2)
    # assert(len(rows) == 18)
    s = ''
    for i, row in enumerate(rows):
        # skip header row
        if i == 0: continue
        # print(row)
        row_crop = img[row[0]:row[1]]
        row_border = cv.copyMakeBorder(row_crop, 1, 1, 1, 1, cv.BORDER_CONSTANT,
                                       value=[0xFF, 0xFF, 0xFF])
        tes.SetImage(Image.fromarray(row_border))
        text = tes.GetUTF8Text().strip()
        # print(text.strip())
        cs = text
        # hack to fix spacing on "SUM" rows
        if i == len(rows)-1 and len(cs) > 4:
            cs = f'{cs[0:3]} {cs[3:]}'
        #global imagenum
        #cv.imwrite(f'row_{imagenum}.png', row_border)
        #imagenum += 1
        # plt.imshow(row_crop, cmap='gray', vmin=0, vmax=255)
        # plt.xticks([]), plt.yticks([])
        # we process each column locally to account for subtle rotation
        # doing this locally for each row should reduce the potential error
        # cs = ''
        # plt.clf()
        for j, col in enumerate(cols):
            break
            if i == 17 and j < 4:
                cs += 'SUM '[j]
                if j == 3:
                    cs += ' '
                continue
            if j == 36:
                cs += ':'
                # don't both
                continue
            # print(i)
            col_crop = img[row[0]-2:row[1]+2, col[0]-2:col[1]+2]
            border = cv.copyMakeBorder(col_crop, 1, 1, 1, 1, cv.BORDER_CONSTANT,
                                       value=[0xFF, 0xFF, 0xFF])
            # global imagenum
            # cv.imwrite(f'sample_{imagenum}.png', border)
            # imagenum += 1
            # sys.exit()
            # plt.subplot(1, len(cols), j + 1)
            # plt.imshow(border, cmap='gray', vmin=0, vmax=255)
            # plt.xticks([]), plt.yticks([])
            tes.SetImage(Image.fromarray(border))
            text = tes.GetUTF8Text()
            text = text.strip()
            if len(text) == 0:
                text = '0'
            text = text.strip()
            # text = identify_char(border)
            # print(text)
            cs += text
            if 2 < j < 36 and j%2 == 1:
                cs += ' '
        print(cs)
        # plt.show()
        s += cs + '\n'
    return s

outfile = open(sys.argv[2], 'w') if len(sys.argv) == 3 else None

# raw image
img = cv.imread(sys.argv[1])

# map to grayscale
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

# first threshold pass
_, thresh1 = cv.threshold(gray, 127, 255, cv.THRESH_BINARY)

# rotate the image to align for later stages
"""
contours, _ = cv.findContours(thresh1, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)
angles = []
for c in contours:
    vx, vy, _, _ = cv.fitLine(c, cv.DIST_L2, 0, 0.01, 0.01)
    a = (180/np.pi)*math.atan2(vy[0], vx[0])
    angles.append(a)
angles = np.array(angles)

lo, hi = np.percentile(angles, (40, 60))
mean = np.mean(angles[np.where((lo <= angles) & (angles <= hi))])
m = cv.getRotationMatrix2D((thresh1.shape[1]//2, thresh1.shape[0]//2), -mean, 1)
rotated = cv.warpAffine(thresh1, m, (thresh1.shape[1], thresh1.shape[0]),
                        cv.INTER_CUBIC)
"""
# blur to create rectangles
blur = cv.GaussianBlur(thresh1,(71,71),0)

# threshold again to add definition
ret,thresh2 = cv.threshold(blur, 253, 255, cv.THRESH_BINARY_INV)


# convert edges to contours and find bounding rectangles
contours, _ = cv.findContours(thresh2, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
# drawing = np.zeros((thresh2.shape[0], thresh2.shape[1]), dtype=np.uint8)
# cv.drawContours(drawing, contours, -1, (255, 255, 255), 3)
# plt.imshow(drawing, cmap='gray', vmin=0, vmax=255)
# plt.show()

boundrect = []
# drawing = np.zeros((thresh2.shape[0], thresh2.shape[1]), dtype=np.uint8)
print(f'Found {len(contours)} contours')
for i, c in enumerate(contours):
    poly = cv.approxPolyDP(c, 3, True)
    rect = tuple(map(int, cv.boundingRect(poly)))
    x, y, w, h = rect
    # filter out noise
    if w*h < 0.01*thresh2.shape[0]*thresh2.shape[1]:
        continue
    # cv.rectangle(drawing, (x, y), (x + w, y + h), (255, 255, 255), 3)
    boundrect.append(rect)

# plt.imshow(drawing, cmap='gray', vmin=0, vmax=255)
# plt.show()

# find median rectangle dimensions
med_w = statistics.median(x[2] for x in boundrect)
med_h = statistics.median(x[3] for x in boundrect)

TOL_W = 0.1
TOL_H = 0.1

min_w = med_w - med_w*0.1
max_w = med_w + med_w*0.1

# height is looser as the last box can be short
min_h = med_h/2 - med_h*0.1
max_h = med_h + med_h*0.1

real_boxes = []
for i, (x, y, w, h) in enumerate(boundrect):
    # reject boxes outside our tolerance
    if min_w <= w <= max_w and min_h <= h <= max_h:
        real_boxes.append((x, y, w, h))

print(f'Reduced to {len(real_boxes)} boxes')

# sort column-wise
real_boxes.sort(key=lambda r: (r[0]/1000)*100000 + r[1])

for x, y, w, h in real_boxes:
    crop = gray[y:y+h, x:x+w]
    s = process_box(crop)
    if outfile:
        outfile.write(s + '\n')
        outfile.flush()

tes.End()
if outfile:
    outfile.close()