create_interest_areas_from_image.py

from PIL import Image, ImageDraw
import pandas as pd
import io
import csv
import os
from pathlib import Path

if os.environ.get('TESSDATA_PREFIX') is None and os.name == 'nt':
    tessdata_prefix = 'C:/Program Files/Tesseract-OCR/tessdata/'
    if Path(tessdata_prefix).exists():
        os.environ['TESSDATA_PREFIX'] = 'C:/Program Files/Tesseract-OCR/tessdata/'
    else:
        tessdata_prefix = None
if os.environ.get('TESSDATA_PREFIX') is None and os.name != 'nt':
    tessdata_prefix = '/usr/share/tesseract-ocr/4.00/tessdata'
    if Path(tessdata_prefix).exists():
        os.environ['TESSDATA_PREFIX'] = '/usr/share/tesseract-ocr/4.00/tessdata'
    else:
        tessdata_prefix = None
    
import pytesseract
if os.name == 'nt':
    if Path(r'c:/Program Files/Tesseract-OCR/tesseract.exe').exists():
        pytesseract.pytesseract.tesseract_cmd = r'c:/Program Files/Tesseract-OCR/tesseract.exe'
else:
    if Path(r'/usr/bin/tesseract').exists():
        pytesseract.pytesseract.tesseract_cmd =r'/usr/bin/tesseract'

def recognize_text(image_path, tesseract_config='--psm 6 -l spa'):
    """
    Performs OCR on an image and returns a DataFrame with character bounding boxes
    and associated information.

    Args:
        image_path: Path to the image file.
        tesseract_config: Configuration string for pytesseract (e.g., '--psm 6 -l spa').

    Returns:
        pandas.DataFrame: DataFrame containing character-level data (df_word_chars).
    """
    # if os.environ['TESSDATA_PREFIX'] is not None:
    #     tesseract_config =  f'--tessdata-dir "{tessdata_prefix}"' + tesseract_config
    image = Image.open(image_path).convert('RGB')
    if hasattr(image_path,'name'):
        im_name = image_path.name
    else:
        im_name = image_path
    image_height = image.height

    # Extract filename for trial_id
    trial_id = os.path.splitext(os.path.basename(im_name))[0]

    # Use pytesseract to extract data for words and characters
    data_words = pytesseract.image_to_data(image, config=tesseract_config)
    data_chars = pytesseract.image_to_boxes(image, config=tesseract_config)

    df_words = pd.read_csv(io.StringIO(data_words), sep='\t', quoting=csv.QUOTE_NONE)
    df_chars = pd.read_csv(io.StringIO(data_chars), sep=' ', header=None, names=['char', 'left', 'top', 'right', 'bottom', 'unknown'])

    # Fix character coordinates
    for index, row in df_chars.iterrows():
        original_top = int(row['top'])
        original_bottom = int(row['bottom'])
        df_chars.at[index, 'top'] = image_height - original_bottom
        df_chars.at[index, 'bottom'] = image_height - original_top

    # Create DataFrame to store spaces
    df_spaces = pd.DataFrame(columns=['level', 'page_num', 'block_num', 'par_num', 'line_num', 'word_num', 'left', 'top', 'width', 'height', 'conf', 'text'])

    # Group words by line, block, and paragraph
    grouped_lines = df_words.groupby(['block_num', 'par_num', 'line_num'])

    for (block_num, par_num, line_num), line_words_df in grouped_lines:
        sorted_words = line_words_df.sort_values(by='left')
        previous_word = None
        for index, current_word in sorted_words.iterrows():
            if previous_word is not None:
                space_left = int(previous_word['left']) + int(previous_word['width'])
                space_width = int(current_word['left']) - space_left
                if space_width > 0:
                    space_top = int(previous_word['top'])
                    space_height = int(previous_word['height'])
                    space_data = {
                        'level': 5,
                        'page_num': int(current_word['page_num']),
                        'block_num': int(current_word['block_num']),
                        'par_num': int(current_word['par_num']),
                        'line_num': int(current_word['line_num']),
                        'word_num': int(previous_word['word_num']),
                        'left': space_left,
                        'top': space_top,
                        'width': space_width,
                        'height': space_height,
                        'conf': 0,
                        'text': ' '
                    }
                    df_spaces = pd.concat([df_spaces, pd.DataFrame(space_data, index=[0])], ignore_index=True)
            previous_word = current_word

    # Create DataFrame for characters within words (and spaces)
    df_word_chars = pd.DataFrame(columns=['char', 'char_xmin', 'char_ymin', 'char_xmax', 'char_ymax',
                                            'block', 'paragraph', 'line_number',
                                            'word_nr', 'letter_nr', 'word',
                                            'char_x_center', 'char_y_center', 'assigned_line', 'trial_id'])


    for index_word, row_word in df_words.iterrows():
        if isinstance(row_word['text'], str) and row_word['text'].strip() and row_word['level'] == 5:
            word_left = int(row_word['left'])
            word_top = int(row_word['top'])
            word_width = int(row_word['width'])
            word_height = int(row_word['height'])
            word_right = word_left + word_width
            word_bottom = word_top + word_height
            word_text = row_word['text']

            char_index_in_word = 0
            relevant_chars = df_chars[
                (df_chars['left'] >= word_left) & (df_chars['right'] <= word_right) &
                (df_chars['top'] >= word_top) & (df_chars['bottom'] <= word_bottom)
            ]
            relevant_chars = relevant_chars.sort_values(by='left')
            previous_char_right = word_left

            for index_char, row_char in relevant_chars.iterrows():
                char_text = row_char['char']
                char_left = previous_char_right
                char_right = int(row_char['right'])
                char_right = min(char_right, word_right)
                if char_left > char_right:
                    char_right = int(row_char['right'])
                char_top = word_top
                char_bottom = word_bottom

                char_data = {
                    'char': char_text,
                    'char_xmin': int(round(char_left)),  # Round and convert to int
                    'char_ymin': int(round(char_top)),   # Round and convert to int
                    'char_xmax': int(round(char_right)),  # Round and convert to int
                    'char_ymax': int(round(char_bottom)), # Round and convert to int
                    'block': int(row_word['block_num']),
                    'paragraph': int(row_word['par_num']),
                    'line_number': int(row_word['line_num']),
                    'word_nr': int(row_word['word_num']),
                    'letter_nr': int(char_index_in_word), #already an int
                    'word': word_text,
                    'char_x_center': int(round((char_left + char_right) / 2)),  # Round and convert
                    'char_y_center': int(round((char_top + char_bottom) / 2)),  # Round and convert
                    'assigned_line': None,
                    'trial_id': trial_id
                }
                df_word_chars = pd.concat([df_word_chars, pd.DataFrame(char_data, index=[0])], ignore_index=True)
                char_index_in_word += 1
                previous_char_right = char_right

            spaces_following_word = df_spaces[
                (df_spaces['word_num'] == int(row_word['word_num'])) &
                (df_spaces['line_num'] == int(row_word['line_num'])) &
                (df_spaces['block_num'] == int(row_word['block_num'])) &
                (df_spaces['par_num'] == int(row_word['par_num']))
            ]

            for index_space, row_space in spaces_following_word.iterrows():
                space_data = {
                    'char': ' ',
                    'char_xmin': int(round(row_space['left'])),    # Round and convert
                    'char_ymin': int(round(row_space['top'])),     # Round and convert
                    'char_xmax': int(round(row_space['left'] + row_space['width'])),  # Round and convert
                    'char_ymax': int(round(row_space['top'] + row_space['height'])), # Round and convert
                    'block': int(row_space['block_num']),
                    'paragraph': int(row_space['par_num']),
                    'line_number': int(row_space['line_num']),
                    'word_nr': int(row_space['word_num']),
                    'letter_nr': int(char_index_in_word),  # Already int
                    'word': word_text,
                    'char_x_center': int(round((row_space['left'] + row_space['left'] + row_space['width']) / 2)), # Round
                    'char_y_center': int(round((row_space['top'] + row_space['top'] + row_space['height']) / 2)), # Round
                    'assigned_line': None,
                    'trial_id': trial_id
                }
                df_word_chars = pd.concat([df_word_chars, pd.DataFrame(space_data, index=[0])], ignore_index=True)
                char_index_in_word += 1

    # Create 'assigned_line' column
    df_word_chars['assigned_line'] = 0
    line_counter = 1
    for block_num in sorted(df_word_chars['block'].unique()):
        for par_num in sorted(df_word_chars.loc[df_word_chars['block'] == block_num, 'paragraph'].unique()):
            for line_num in sorted(df_word_chars.loc[(df_word_chars['block'] == block_num) & (df_word_chars['paragraph'] == par_num), 'line_number'].unique()):
                line_mask = (df_word_chars['line_number'] == line_num) & (df_word_chars['paragraph'] == par_num) & (df_word_chars['block'] == block_num)
                df_word_chars.loc[line_mask, 'assigned_line'] = line_counter
                line_counter += 1

    # Adjust Y_Start, Y_End, and char_y_center, converting to integers
    for assigned_line in df_word_chars['assigned_line'].unique():
        line_mask = (df_word_chars['assigned_line'] == assigned_line)
        min_top = df_word_chars.loc[line_mask, 'char_ymin'].min()
        max_bottom = df_word_chars.loc[line_mask, 'char_ymax'].max()
        new_y_center = (min_top + max_bottom) / 2
        df_word_chars.loc[line_mask, 'char_ymin'] = int(round(min_top))      # Round and convert
        df_word_chars.loc[line_mask, 'char_ymax'] = int(round(max_bottom))     # Round and convert
        df_word_chars.loc[line_mask, 'char_y_center'] = int(round(new_y_center)) # Round and convert

    # Convert relevant columns to integers
    int_columns = ['char_xmin', 'char_ymin', 'char_xmax', 'char_ymax', 'block', 'paragraph',
                    'line_number', 'word_nr', 'letter_nr', 'char_x_center', 'char_y_center', 'assigned_line']
    for col in int_columns:
        df_word_chars[col] = df_word_chars[col].astype(int)

    return df_word_chars


def draw_char_boxes(image_path, df_word_chars, output_path='output_boxes_combined.png'):
    """
    Draws bounding boxes around characters on the image.

    Args:
        image_path: Path to the image file.
        df_word_chars: DataFrame containing character bounding box data.
        output_path: Path to save the image with bounding boxes.  Defaults to 'output_boxes_combined.png'.
    """
    image = Image.open(image_path).convert('RGB')
    draw = ImageDraw.Draw(image)

    # Draw bounding boxes for characters (purple)
    for index, row in df_word_chars.iterrows():
        left = int(row['char_xmin'])
        top = int(row['char_ymin'])
        right = int(row['char_xmax'])
        bottom = int(row['char_ymax'])
        draw.rectangle([(left, top), (right, bottom)], outline='purple', width=1)

    # Display or save the image
    image.save(output_path)


# Example usage
if __name__ == '__main__':
    # image_path = 'testfiles/testim_ocr.png'
    image_path = 'testfiles/newplot.png'
    # Example with default tesseract config
    df_chars = recognize_text(image_path)
    draw_char_boxes(image_path, df_chars)
    df_chars.to_csv('testim_ocr_df_word_chars_test.csv', index=False)
    print("\nDataFrame of Characters within Words (df_word_chars) - Default Config:")
    print(df_chars)