from PIL import Image, ImageEnhance, ImageOps from textractor import Textractor from textractor.visualizers.entitylist import EntityList from textractor.data.constants import TextractFeatures from openpyxl import load_workbook from openpyxl.utils.cell import range_boundaries

def enhance_image(img):

   enhancer = ImageEnhance.Contrast(img)
   enhanced_img = enhancer.enhance(1.5)
   return enhanced_img

def ocr(input_file, output_dir, image=None):

   extractor = Textractor(region_name="us-east-1")
   output_filename = os.path.splitext(input_file)[0] + '.xlsx'
   output_filepath = os.path.join(output_dir, output_filename)
   if image is None:
       print('input:', input_file)
       image = Image.open(input_file)
   image = enhance_image(image)
   document = extractor.analyze_document(
       file_source=image,
       features=[TextractFeatures.TABLES],
       save_image=True
   )
   table = EntityList(document.tables[0])
   table[0].to_excel(filepath=output_filepath)
   wb = load_workbook(filename=output_filepath)
   for st_name in wb.sheetnames:
       st = wb[st_name]
       mcr_coord_list = [mcr.coord for mcr in st.merged_cells.ranges]
       for mcr in mcr_coord_list:
           min_col, min_row, max_col, max_row = range_boundaries(mcr)
           top_left_cell_value = st.cell(row=min_row, column=min_col).value
           st.unmerge_cells(mcr)
           for row in st.iter_rows(min_col=min_col, min_row=min_row, max_col=max_col, max_row=max_row):
               for cell in row:
                   cell.value = top_left_cell_value
   wb.save(output_filepath)
   return output_filepath