Code clean up and image orientation function fix suggested in PR

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 = [], []

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")