From 45e73eab0fefbf2976a9e6f81bd16216d5b1e2b5 Mon Sep 17 00:00:00 2001 From: anjus Date: Fri, 9 Aug 2024 09:52:45 -0400 Subject: [PATCH] Code clean up and image orientation function fix suggested in PR --- app_doctr.py | 15 ++++++-------- src/msfocr/doctr/ocr_functions.py | 34 +++++++++++++++---------------- 2 files changed, 23 insertions(+), 26 deletions(-) diff --git a/app_doctr.py b/app_doctr.py index d7b1c7f..7a19b4a 100644 --- a/app_doctr.py +++ b/app_doctr.py @@ -50,7 +50,7 @@ def create_ocr(): doctr_ocr = DocTR(detect_language=False) return ocr_model, doctr_ocr -@st.cache_data +@st.cache_data(show_spinner=False) def get_uploaded_images(tally_sheet): """ List of images uploaded by user as docTR DocumentFiles @@ -64,7 +64,7 @@ def get_uploaded_images(tally_sheet): res.append(DocumentFile.from_images(image)) return res -@st.cache_data +@st.cache_data(show_spinner=False) def get_results(uploaded_images): return [doctr_ocr_functions.get_word_level_content(ocr_model, doc) for doc in uploaded_images] @@ -297,15 +297,11 @@ def clean_up(table_dfs): _List_: List of table data frames """ for table in table_dfs: - print(table) for row in range(table.shape[0]): for col in range(table.shape[1]): cell_value = table.iloc[row][col] - print(cell_value) if cell_value is None or cell_value=="None": - table.iloc[row][col] = "" - print(table.iloc[row][col]) - print(table_dfs) + table.iloc[row][col] = "" return table_dfs # Initializing session state variables that only need to be set on startup @@ -455,8 +451,9 @@ def authenticate(): # *************************************** # Populate streamlit with data recognized from tally sheets - uploaded_images = get_uploaded_images(tally_sheet_images) - results = get_results(uploaded_images) + with st.spinner("Running image recognition..."): + uploaded_images = get_uploaded_images(tally_sheet_images) + results = get_results(uploaded_images) # Spinner for data upload. If it's going to be on screen for long, make it bespoke table_dfs, page_nums_to_display = [], [] diff --git a/src/msfocr/doctr/ocr_functions.py b/src/msfocr/doctr/ocr_functions.py index 0fd89ad..33fa5a1 100644 --- a/src/msfocr/doctr/ocr_functions.py +++ b/src/msfocr/doctr/ocr_functions.py @@ -114,7 +114,6 @@ def get_tabular_content(model, image, confidence_dict=None): return table_df, confidence_df - def get_sheet_type(res): """ Finds the type of the tally sheet (dataSet, orgUnit, period) from the result of OCR model, where @@ -208,22 +207,23 @@ def correct_image_orientation(image_path): :param image_path: The path to the image file. :return: PIL.Image.Image: The image with corrected orientation. """ - with Image.open(image_path) as image: - orientation = None - try: - for orientation in ExifTags.TAGS.keys(): - if ExifTags.TAGS[orientation] == 'Orientation': - break - exif = dict(image.getexif().items()) - if exif.get(orientation) == 3: - image = image.rotate(180, expand=True) - elif exif.get(orientation) == 6: - image = image.rotate(270, expand=True) - elif exif.get(orientation) == 8: - image = image.rotate(90, expand=True) - except (AttributeError, KeyError, IndexError): - pass - return image.copy() + with Image.open(image_path) as image: + image.load() + orientation = None + try: + for orientation in ExifTags.TAGS.keys(): + if ExifTags.TAGS[orientation] == 'Orientation': + break + exif = dict(image.getexif().items()) + if exif.get(orientation) == 3: + image = image.rotate(180, expand=True) + elif exif.get(orientation) == 6: + image = image.rotate(270, expand=True) + elif exif.get(orientation) == 8: + image = image.rotate(90, expand=True) + except (AttributeError, KeyError, IndexError): + pass + return image # ocr_model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True) # document = DocumentFile.from_images("IMG_20240514_090947.jpg")