generate_tesseract_results.py

from pytesseract import Output
import pytesseract
from glob import glob
import cv2
from tqdm import tqdm
import os
import json


def get_tesseract_results(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    ocr_result = pytesseract.image_to_data(
        image, config="--oem 1 --psm 3", output_type=Output.DICT
    )
    return ocr_result


# TODO: filter out ocr results that don't have a corresponding xml.
if __name__ == "__main__":

    dataset_dir = "./assets/images/mantis/labeled"
    images_dir = "./assets/images/mantis"
    tesseract_results = os.path.join(dataset_dir, "tesseract_results_lstm")

    assert os.path.exists(images_dir), "images directory doesn't exist"
    if not os.path.exists(tesseract_results):
        os.makedirs(tesseract_results)

    images = glob(os.path.join(images_dir, "*.jpg"))

    for image in tqdm(images, desc="Generating Tesseract Results"):
        image_name = os.path.splitext(os.path.split(image)[-1])[0]
        result = get_tesseract_results(image)
        with open(os.path.join(tesseract_results, image_name + ".json"), "w") as f:
            json.dump(result, f)