chore: import upstream snapshot with attribution

test_unstructured/partition/utils/ocr_models/test_ocr_interface.py

@@ -0,0 +1,114 @@
+# pyright: reportPrivateUsage=false
+
+"""Unit-test suite for the `unstructured.partition.utils.ocr_models.ocr_interface` module."""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from test_unstructured.unit_utils import (
+    FixtureRequest,
+    LogCaptureFixture,
+    Mock,
+    instance_mock,
+    method_mock,
+    property_mock,
+)
+from unstructured.partition.utils.config import ENVConfig
+from unstructured.partition.utils.constants import (
+    OCR_AGENT_PADDLE,
+    OCR_AGENT_PADDLE_OLD,
+    OCR_AGENT_TESSERACT,
+    OCR_AGENT_TESSERACT_OLD,
+)
+from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
+
+
+class DescribeOCRAgent:
+    """Unit-test suite for `unstructured.partition.utils...ocr_interface.OCRAgent` class."""
+
+    def it_provides_access_to_the_configured_OCR_agent(
+        self, _get_ocr_agent_cls_qname_: Mock, get_instance_: Mock, ocr_agent_: Mock
+    ):
+        _get_ocr_agent_cls_qname_.return_value = OCR_AGENT_TESSERACT
+        get_instance_.return_value = ocr_agent_
+
+        ocr_agent = OCRAgent.get_agent(language="eng")
+
+        _get_ocr_agent_cls_qname_.assert_called_once_with()
+        get_instance_.assert_called_once_with(OCR_AGENT_TESSERACT, "eng")
+        assert ocr_agent is ocr_agent_
+
+    def but_it_raises_when_the_requested_agent_is_not_whitelisted(
+        self, _get_ocr_agent_cls_qname_: Mock
+    ):
+        _get_ocr_agent_cls_qname_.return_value = "Invalid.Ocr.Agent.Qname"
+        with pytest.raises(ValueError, match="must be set to a whitelisted module"):
+            OCRAgent.get_agent(language="eng")
+
+    @pytest.mark.parametrize("exception_cls", [ImportError, AttributeError])
+    def and_it_raises_when_the_requested_agent_cannot_be_loaded(
+        self, _get_ocr_agent_cls_qname_: Mock, exception_cls: type[Exception], _clear_cache
+    ):
+        _get_ocr_agent_cls_qname_.return_value = OCR_AGENT_TESSERACT
+        with (
+            patch(
+                "unstructured.partition.utils.ocr_models.ocr_interface.importlib.import_module",
+                side_effect=exception_cls,
+            ),
+            pytest.raises(RuntimeError, match="Could not get the OCRAgent instance"),
+        ):
+            OCRAgent.get_agent(language="eng")
+
+    @pytest.mark.parametrize(
+        ("OCR_AGENT", "expected_value"),
+        [
+            (OCR_AGENT_PADDLE, OCR_AGENT_PADDLE),
+            (OCR_AGENT_PADDLE_OLD, OCR_AGENT_PADDLE),
+            (OCR_AGENT_TESSERACT, OCR_AGENT_TESSERACT),
+            (OCR_AGENT_TESSERACT_OLD, OCR_AGENT_TESSERACT),
+        ],
+    )
+    def it_computes_the_OCR_agent_qualified_module_name(
+        self, OCR_AGENT: str, expected_value: str, OCR_AGENT_prop_: Mock
+    ):
+        OCR_AGENT_prop_.return_value = OCR_AGENT
+        assert OCRAgent._get_ocr_agent_cls_qname() == expected_value
+
+    @pytest.mark.parametrize("OCR_AGENT", [OCR_AGENT_PADDLE_OLD, OCR_AGENT_TESSERACT_OLD])
+    def and_it_logs_a_warning_when_the_OCR_AGENT_module_name_is_obsolete(
+        self, caplog: LogCaptureFixture, OCR_AGENT: str, OCR_AGENT_prop_: Mock
+    ):
+        OCR_AGENT_prop_.return_value = OCR_AGENT
+        OCRAgent._get_ocr_agent_cls_qname()
+        assert f"OCR agent name {OCR_AGENT} is outdated " in caplog.text
+
+    # -- fixtures --------------------------------------------------------------------------------
+
+    @pytest.fixture()
+    def _clear_cache(self):
+        # Clear the cache created by @functools.lru_cache(maxsize=None) on OCRAgent.get_instance()
+        # before each test
+        OCRAgent.get_instance.cache_clear()
+        yield
+        # Clear the cache created by @functools.lru_cache(maxsize=None) on OCRAgent.get_instance()
+        # after each test (just in case)
+        OCRAgent.get_instance.cache_clear()
+
+    @pytest.fixture()
+    def get_instance_(self, request: FixtureRequest):
+        return method_mock(request, OCRAgent, "get_instance")
+
+    @pytest.fixture()
+    def _get_ocr_agent_cls_qname_(self, request: FixtureRequest):
+        return method_mock(request, OCRAgent, "_get_ocr_agent_cls_qname")
+
+    @pytest.fixture()
+    def ocr_agent_(self, request: FixtureRequest):
+        return instance_mock(request, OCRAgent)
+
+    @pytest.fixture()
+    def OCR_AGENT_prop_(self, request: FixtureRequest):
+        return property_mock(request, ENVConfig, "OCR_AGENT")

test_unstructured/partition/utils/test_config.py

@@ -0,0 +1,58 @@
+import shutil
+import tempfile
+from pathlib import Path
+
+import pytest
+
+
+def test_default_config():
+    from unstructured.partition.utils.config import env_config
+
+    assert env_config.IMAGE_CROP_PAD == 0
+
+
+def test_env_override(monkeypatch):
+    monkeypatch.setenv("IMAGE_CROP_PAD", str(1))
+    from unstructured.partition.utils.config import env_config
+
+    assert env_config.IMAGE_CROP_PAD == 1
+
+
+@pytest.fixture()
+def _setup_tmpdir():
+    from unstructured.partition.utils.config import env_config
+
+    _tmpdir = tempfile.tempdir
+    _storage_tmpdir = env_config.GLOBAL_WORKING_PROCESS_DIR
+    _storage_tmpdir_bak = f"{env_config.GLOBAL_WORKING_PROCESS_DIR}_bak"
+    if Path(_storage_tmpdir).is_dir():
+        shutil.move(_storage_tmpdir, _storage_tmpdir_bak)
+        tempfile.tempdir = None
+    yield
+    if Path(_storage_tmpdir_bak).is_dir():
+        if Path(_storage_tmpdir).is_dir():
+            shutil.rmtree(_storage_tmpdir)
+        shutil.move(_storage_tmpdir_bak, _storage_tmpdir)
+        tempfile.tempdir = _tmpdir
+
+
+@pytest.mark.usefixtures("_setup_tmpdir")
+def test_env_storage_disabled(monkeypatch):
+    monkeypatch.setenv("GLOBAL_WORKING_DIR_ENABLED", "false")
+    from unstructured.partition.utils.config import env_config
+
+    assert not env_config.GLOBAL_WORKING_DIR_ENABLED
+    assert str(Path.home() / ".cache/unstructured") == env_config.GLOBAL_WORKING_DIR
+    assert not Path(env_config.GLOBAL_WORKING_PROCESS_DIR).is_dir()
+    assert tempfile.gettempdir() != env_config.GLOBAL_WORKING_PROCESS_DIR
+
+
+@pytest.mark.usefixtures("_setup_tmpdir")
+def test_env_storage_enabled(monkeypatch):
+    monkeypatch.setenv("GLOBAL_WORKING_DIR_ENABLED", "true")
+    from unstructured.partition.utils.config import env_config
+
+    assert env_config.GLOBAL_WORKING_DIR_ENABLED
+    assert str(Path.home() / ".cache/unstructured") == env_config.GLOBAL_WORKING_DIR
+    assert Path(env_config.GLOBAL_WORKING_PROCESS_DIR).is_dir()
+    assert tempfile.gettempdir() == env_config.GLOBAL_WORKING_PROCESS_DIR

test_unstructured/partition/utils/test_sorting.py

@@ -0,0 +1,157 @@
+import numpy as np
+import pytest
+from unstructured_inference.inference.elements import TextRegions
+
+from unstructured.documents.coordinates import PixelSpace
+from unstructured.documents.elements import CoordinatesMetadata, Element, Text
+from unstructured.partition.utils.constants import SORT_MODE_BASIC, SORT_MODE_XY_CUT
+from unstructured.partition.utils.sorting import (
+    coord_has_valid_points,
+    coordinates_to_bbox,
+    shrink_bbox,
+    sort_page_elements,
+    sort_text_regions,
+)
+
+
+class MockCoordinatesMetadata(CoordinatesMetadata):
+    def __init__(self, points):
+        system = PixelSpace(width=300, height=500)
+
+        super().__init__(points, system)
+
+
+def test_coord_valid_coordinates():
+    coordinates = CoordinatesMetadata([(1, 2), (3, 4), (5, 6), (7, 8)], PixelSpace)
+    assert coord_has_valid_points(coordinates) is True
+
+
+def test_coord_missing_incomplete_point():
+    coordinates = CoordinatesMetadata([(1, 2), (3, 4), (5, 6)], PixelSpace)
+    assert coord_has_valid_points(coordinates) is False
+
+
+def test_coord_negative_values():
+    coordinates = CoordinatesMetadata([(1, 2), (3, 4), (5, -6), (7, 8)], PixelSpace)
+    assert coord_has_valid_points(coordinates) is False
+
+
+def test_coord_weird_values():
+    coordinates = CoordinatesMetadata([(1, 2), ("3", 4), (5, 6), (7, 8)], PixelSpace)
+    assert coord_has_valid_points(coordinates) is False
+
+
+def test_coord_invalid_point_structure():
+    coordinates = CoordinatesMetadata([(1, 2), (3, 4, 5), (6, 7), (8, 9)], PixelSpace)
+    assert coord_has_valid_points(coordinates) is False
+
+
+@pytest.mark.parametrize("sort_mode", ["xy-cut", "basic"])
+def test_sort_page_elements_without_coordinates(sort_mode):
+    elements = [Element(str(idx)) for idx in range(5)]
+    assert sort_page_elements(elements) == elements
+
+
+def test_sort_xycut_neg_coordinates():
+    elements = []
+    for idx in range(2):
+        elem = Text(str(idx))
+        elem.metadata.coordinates = CoordinatesMetadata(
+            [(0, idx), (3, 4), (6, 7), (8, 9)],
+            PixelSpace,
+        )
+        elements.append(elem)
+
+    # NOTE(crag): xycut not attempted, sort_page_elements returns original list
+    assert sort_page_elements(elements, sort_mode=SORT_MODE_XY_CUT) is not elements
+
+
+def test_sort_xycut_pos_coordinates():
+    elements = []
+    for idx in range(2):
+        elem = Text(str(idx))
+        elem.metadata.coordinates = CoordinatesMetadata(
+            [(1, 2), (3, 4), (6, 7), (8, 9)],
+            PixelSpace,
+        )
+        elements.append(elem)
+
+    # NOTE(crag): xycut ran, so different list reference returned from input list
+    assert sort_page_elements(elements, sort_mode=SORT_MODE_XY_CUT) is not elements
+
+
+def test_sort_basic_neg_coordinates():
+    elements = []
+    for idx in range(3):
+        elem = Text(str(idx))
+        elem.metadata.coordinates = CoordinatesMetadata(
+            [(1, -idx), (3, 4), (6, 7), (8, 9)],
+            PixelSpace,
+        )
+        elements.append(elem)
+
+    sorted_page_elements = sort_page_elements(elements, sort_mode=SORT_MODE_BASIC)
+    sorted_elem_text = " ".join([str(elem.text) for elem in sorted_page_elements])
+    assert sorted_elem_text == "2 1 0"
+
+
+def test_sort_basic_pos_coordinates():
+    elements = []
+    for idx in range(3):
+        elem = Text(str(9 - idx))
+        elem.metadata.coordinates = CoordinatesMetadata(
+            [(1, 9 - idx), (3, 4), (6, 7), (8, 9)],
+            PixelSpace,
+        )
+        elements.append(elem)
+
+    sorted_page_elements = sort_page_elements(elements, sort_mode=SORT_MODE_BASIC)
+    assert sorted_page_elements is not elements
+
+    sorted_elem_text = " ".join([str(elem.text) for elem in sorted_page_elements])
+    assert sorted_elem_text == "7 8 9"
+
+
+def test_sort_text_regions():
+    unsorted = TextRegions(
+        element_coords=np.array(
+            [[1, 2, 2, 2], [1, 1, 2, 2], [3, 1, 4, 4]],
+        ),
+        texts=np.array(["1", "2", "3"]),
+        sources=np.array(["foo"] * 3),
+    )
+    assert sort_text_regions(unsorted, sort_mode=SORT_MODE_BASIC).texts.tolist() == ["2", "3", "1"]
+
+
+@pytest.mark.parametrize(
+    "coords",
+    [
+        [[1, 2, 2, 2], [1, 1, 2, 2], [3, -1, 4, 4]],
+        [[1, 2, 2, 2], [1, 1, 2, 2], [3, None, 4, 4]],
+    ],
+)
+def test_sort_text_regions_with_invalid_coords_using_xy_cut_does_no_ops(coords):
+    unsorted = TextRegions(
+        element_coords=np.array(coords).astype(float),
+        texts=np.array(["1", "2", "3"]),
+        sources=np.array(["foo"] * 3),
+    )
+    assert sort_text_regions(unsorted).texts.tolist() == ["1", "2", "3"]
+
+
+def test_coordinates_to_bbox():
+    coordinates_data = MockCoordinatesMetadata([(10, 20), (10, 200), (100, 200), (100, 20)])
+    expected_result = (10, 20, 100, 200)
+    assert coordinates_to_bbox(coordinates_data) == expected_result
+
+
+def test_shrink_bbox():
+    bbox = (0, 0, 200, 100)
+    shrink_factor = 0.9
+    expected_result = (0, 0, 180, 90)
+    assert shrink_bbox(bbox, shrink_factor) == expected_result
+
+    bbox = (20, 20, 320, 120)
+    shrink_factor = 0.9
+    expected_result = (20, 20, 290, 110)
+    assert shrink_bbox(bbox, shrink_factor) == expected_result

test_unstructured/partition/utils/test_xycut.py

@@ -0,0 +1,181 @@
+from unittest.mock import patch
+
+import cv2
+import numpy as np
+import pytest
+
+from unstructured.partition.utils import xycut
+
+
+def test_projection_by_bboxes():
+    boxes = np.array([[10, 20, 50, 60], [30, 40, 70, 80]])
+
+    # Test case 1: Horizontal projection
+    result_horizontal = xycut.projection_by_bboxes(boxes, 0)
+    expected_result_horizontal = np.array(
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+    )
+    assert np.array_equal(result_horizontal[:30], expected_result_horizontal)
+
+    # Test case 2: Vertical projection
+    result_vertical = xycut.projection_by_bboxes(boxes, 1)
+    expected_result_vertical = np.array(
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+    )
+    assert np.array_equal(result_vertical[:30], expected_result_vertical)
+
+
+def test_split_projection_profile():
+    # Test case 1: Sample projection profile with given min_value and min_gap
+    arr_values = np.array([0, 0, 3, 4, 0, 0, 2, 0, 0, 0, 5, 6, 7, 0, 0, 0])
+    min_value = 0
+    min_gap = 1
+    result = xycut.split_projection_profile(arr_values, min_value, min_gap)
+    expected_result = (np.array([2, 6, 10]), np.array([4, 7, 13]))
+    assert np.array_equal(result, expected_result)
+
+    # Test case 2: Another sample projection profile with different parameters
+    arr_values = np.array([0, 2, 0, 0, 0, 3, 0, 0, 4, 5, 6, 0, 0, 0])
+    min_value = 1
+    min_gap = 2
+    result = xycut.split_projection_profile(arr_values, min_value, min_gap)
+    expected_result = (np.array([1, 5, 8]), np.array([2, 6, 11]))
+    assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    ("recursive_func", "expected"),
+    [
+        (xycut.recursive_xy_cut, [0, 1, 2]),
+        (xycut.recursive_xy_cut_swapped, [0, 2, 1]),
+    ],
+)
+def test_recursive_xy_cut(recursive_func, expected):
+    boxes = np.array([[0, 0, 20, 20], [200, 0, 230, 30], [0, 40, 50, 50]])
+    indices = np.array([0, 1, 2])
+    res = []
+    recursive_func(boxes, indices, res)
+    assert res == expected
+
+
+def test_points_to_bbox():
+    # Test a valid case
+    points = [10, 20, 30, 40, 50, 60, 70, 80]
+    result = xycut.points_to_bbox(points)
+    assert result == [10, 20, 70, 80]
+
+    # Test a case where points are unordered
+    points = [30, 40, 10, 20, 70, 80, 50, 60]
+    result = xycut.points_to_bbox(points)
+    assert result == [10, 20, 70, 80]
+
+    # Test a case where all points are negative
+    points = [-10, -20, -30, -40, -50, -60, -70, -80]
+    result = xycut.points_to_bbox(points)
+    assert result == [0, 0, 0, 0]
+
+    # Test a case with invalid number of points
+    with pytest.raises(AssertionError):
+        points = [10, 20, 30, 40, 50, 60]  # Missing two points
+        xycut.points_to_bbox(points)
+
+
+def test_bbox2points():
+    # Test a valid case
+    bbox = [10, 20, 70, 80]
+    result = xycut.bbox2points(bbox)
+    assert result == [10, 20, 70, 20, 70, 80, 10, 80]
+
+    # Test a case where the top and bottom are the same
+    bbox = [10, 20, 70, 20]
+    result = xycut.bbox2points(bbox)
+    assert result == [10, 20, 70, 20, 70, 20, 10, 20]
+
+    # Test a case where left and right are the same
+    bbox = [10, 20, 10, 80]
+    result = xycut.bbox2points(bbox)
+    assert result == [10, 20, 10, 20, 10, 80, 10, 80]
+
+    # Test a case where the bbox is a point (left and right are the same,
+    # top and bottom are the same)
+    bbox = [10, 20, 10, 20]
+    result = xycut.bbox2points(bbox)
+    assert result == [10, 20, 10, 20, 10, 20, 10, 20]
+
+
+def test_vis_polygon():
+    img = np.ones((200, 200, 3), dtype=np.uint8) * 255
+    points = [(50, 50), (150, 50), (150, 150), (50, 150)]
+    color = (0, 0, 255)  # Red color
+    thickness = 2
+
+    result_img = xycut.vis_polygon(img, points, thickness, color)
+
+    # Define the expected image with the square drawn
+    expected_img = np.copy(img)
+    cv2.line(expected_img, points[0], points[1], color, thickness)
+    cv2.line(expected_img, points[1], points[2], color, thickness)
+    cv2.line(expected_img, points[2], points[3], color, thickness)
+    cv2.line(expected_img, points[3], points[0], color, thickness)
+
+    assert np.array_equal(result_img, expected_img)
+
+
+def test_vis_points():
+    img = np.ones((200, 200, 3), dtype=np.uint8) * 255
+    points = [[10, 20, 30, 20, 30, 40, 10, 40], [50, 60, 70, 60, 70, 80, 50, 80]]
+    texts = ["Label1", "Label2"]
+    color = (0, 200, 0)
+
+    result_img = xycut.vis_points(img, points, texts, color)
+
+    # Check if the resulting image contains the expected shapes and labels
+    expected_img = np.copy(img)
+
+    # Draw polygons and labels for each set of points
+    for i, _points in enumerate(points):
+        xycut.vis_polygon(expected_img, np.array(_points).reshape(-1, 2), thickness=2, color=color)
+        bbox = xycut.points_to_bbox(_points)
+        left, top, right, bottom = bbox
+        cx = (left + right) // 2
+        cy = (top + bottom) // 2
+        txt = texts[i]
+
+        # Draw a filled rectangle for the label background
+        cat_size = cv2.getTextSize(txt, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
+        expected_img = cv2.rectangle(
+            expected_img,
+            (cx - 5 * len(txt), cy - cat_size[1] - 5),
+            (cx - 5 * len(txt) + cat_size[0], cy - 5),
+            color,
+            -1,
+        )
+
+        # Draw the label text
+        expected_img = cv2.putText(
+            expected_img,
+            txt,
+            (cx - 5 * len(txt), cy - 5),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (255, 255, 255),
+            thickness=1,
+            lineType=cv2.LINE_AA,
+        )
+
+    assert np.array_equal(result_img, expected_img)
+
+
+def test_vis_polygons_with_index():
+    img = np.ones((200, 200, 3), dtype=np.uint8) * 255
+    points = [[10, 20, 30, 20, 30, 40, 10, 40], [50, 60, 70, 60, 70, 80, 50, 80]]
+
+    with patch(
+        "unstructured.partition.utils.xycut.vis_points", return_value=img
+    ) as mock_vis_points:
+        result_img = xycut.vis_polygons_with_index(img, points)
+
+        # Check if vis_points was called with the correct arguments
+        mock_vis_points.assert_called_once()
+
+        assert np.array_equal(result_img, img)