9. Introducing OCR Example

##############################################################
### Activate conda env and run the script
#     module load anaconda3
#     source activate /zfs/gsb/intermediate-yens/conda/ocr
#     python ocr-serial.py
##############################################################
import numpy as np
import pandas as pd
from PIL import Image
import pytesseract # OCR with Tesseract
import glob, os, shutil, csv, json, time
from os import makedirs
from os.path import exists, isdir
from datetime import datetime
from tqdm import tqdm   # progress bar

# start the timer
tmp = time.time()

# store OCR results
out_path = 'ocr_serial_out/'

# remove output directory if it already exists
if exists(out_path) and isdir(out_path):
    shutil.rmtree(out_path)

# make empty dir for OCR results
makedirs(out_path)

# make a master df -> CSV file to keep track of what we have already processed

# check if file exists, if it doesn't, make a new CSV file
if not exists(out_path + 'processed_images.csv'):
    # create a new csv file
    with open(out_path + 'processed_images.csv', 'w') as csvfile:
        filewriter = csv.writer(csvfile, delimiter=',', quoting=csv.QUOTE_NONE)
        # write headers
        filewriter.writerow(['image_path', 'image_name', 'ocr_text', 'date'])


# csv exists so read it in
df = pd.read_csv(out_path + 'processed_images.csv')

# path to the data
data_path = '/scratch/darc/intermediate-yens/data/'

# make a list of image paths to OCR
image_paths = glob.glob(data_path + '*.jpg')

tot_images = len(image_paths)

# drop images we already processed
image_paths = [img for img in image_paths if img not in set(df['image_path'])]

print('image number to process: ', len(image_paths))

intermediateFileNum = 100
ncore = 1

# number of image batches
if len(image_paths) % intermediateFileNum == 0:
    nbatches = int(len(image_paths) / intermediateFileNum)
else:
    nbatches = int(len(image_paths) / intermediateFileNum) + 1

# make progress bar for the for loop over batches of images
for i in tqdm( range(nbatches) ):
    # i is the batch number
    # write intermediate results while processing all of the files
    # make a temp df to append to result

    df_tmp = pd.DataFrame(columns=['image_path', 'image_name', 'ocr_text', 'date'])
    start = i * intermediateFileNum

    if i == int(len(image_paths) / intermediateFileNum):
        end = len(image_paths)
    else:
        end = i * intermediateFileNum + intermediateFileNum

    image_paths_batch = image_paths[start: end]

    image_names = []
    text_from_ocr = []

    # loop over images to convert them into text in batches
    for img in image_paths_batch:
        text_from_ocr.append( json.dumps(pytesseract.image_to_string(Image.open(img)).strip() ))
        image_names.append(img[32:])

    # make a pd series for image paths
    df_tmp['image_path'] = pd.Series(image_paths_batch, dtype='object')

    # make a pd series for image names
    df_tmp['image_name'] = pd.Series(image_names, dtype='object')

    # make a pd series for OCR results
    df_tmp['ocr_text'] = pd.Series(text_from_ocr, dtype='object')

    # make a pd series for date
    date_list = [datetime.today().strftime("%m/%d/%Y") for i in range(len(image_paths_batch))]
    df_tmp['date'] = pd.Series(date_list, dtype='object')

    # add data to df
    df = df.append(df_tmp)

    # save results
    df.to_csv(out_path + 'processed_images.csv', index = False)
    print('processed', len(df), 'images out of', tot_images)

print('running with %d core took: %s seconds' % (ncore, str(time.time() - tmp)))