Fix empty df scenario

Unstructured-IO · Jan 8, 2025 · 2d9054d · 2d9054d
1 parent 0e44926
commit 2d9054d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/test_unstructured/partition/pdf_image/test_ocr.py b/test_unstructured/partition/pdf_image/test_ocr.py
@@ -551,3 +551,14 @@ def test_hocr_to_dataframe():
     assert df["width"].iloc[0] == 60
     assert df["height"].iloc[0] == 13
     assert df["text"].iloc[0] == "word"
+
+
+def test_hocr_to_dataframe_when_no_prediction_empty_df():
+    df = OCRAgentTesseract().hocr_to_dataframe(hocr="")
+
+    assert df.shape == (0, 5)
+    assert "left" in df.columns
+    assert "top" in df.columns
+    assert "width" in df.columns
+    assert "text" in df.columns
+    assert "text" in df.columns
diff --git a/unstructured/partition/utils/ocr_models/tesseract_ocr.py b/unstructured/partition/utils/ocr_models/tesseract_ocr.py
@@ -132,7 +132,7 @@ def hocr_to_dataframe(
                         "text": text,
                     }
                 )
-        ocr_df = pd.DataFrame(df_entries)
+        ocr_df = pd.DataFrame(df_entries, columns=["left", "top", "width", "height", "text"])
         return ocr_df
 
     @staticmethod