def test_should_filter_by_score(self):
layout = Layout([
TextBlock(Rectangle(11, 10, 100, 100),
text='block1',
type='Test',
score=0.4),
TextBlock(Rectangle(12, 10, 100, 100),
text='block2',
type='Test',
score=0.5),
TextBlock(Rectangle(13, 10, 100, 100),
text='block3',
type='Test',
score=0.6)
])
result = LayoutParserComputerVisionModelResult(layout,
score_threshold=0.5,
avoid_overlapping=False)
instances = result.get_instances_by_type_name('Test')
LOGGER.debug('instances: %r', instances)
bounding_boxes = [
instance.get_bounding_box() for instance in instances
]
assert bounding_boxes == [
BoundingBox(12, 10, 100 - 12, 100 - 10),
BoundingBox(13, 10, 100 - 13, 100 - 10)
]
def test_should_replace_text_and_graphics_within_bounding_box_of_semantic_graphics(
self):
page_coordinates = LayoutPageCoordinates.from_bounding_box(
BoundingBox(0, 0, 200, 200), page_number=1)
semantic_graphic_coordinates = LayoutPageCoordinates.from_bounding_box(
BoundingBox(10, 90, 100, 50), page_number=1)
keep_coordinates = LayoutPageCoordinates.from_bounding_box(
BoundingBox(10, 10, 100, 20), page_number=1)
remove_coordinates = LayoutPageCoordinates.from_bounding_box(
BoundingBox(10, 100, 100, 20), page_number=1)
empty_coordinates = LayoutPageCoordinates.from_bounding_box(
BoundingBox(10, 100, 0, 0), page_number=1)
keep_token = LayoutToken('keep', coordinates=keep_coordinates)
remove_token = LayoutToken('remove', coordinates=remove_coordinates)
keep_graphic = LayoutGraphic(coordinates=keep_coordinates,
graphic_type='keep-graphic')
remove_graphic = LayoutGraphic(coordinates=remove_coordinates,
graphic_type='remove-graphic')
layout_document = LayoutDocument(pages=[
LayoutPage(blocks=[
LayoutBlock(
lines=[LayoutLine(tokens=[keep_token, remove_token])])
],
graphics=[keep_graphic, remove_graphic],
meta=LayoutPageMeta(
page_number=page_coordinates.page_number,
coordinates=page_coordinates))
])
layout_graphic = LayoutGraphic(
coordinates=semantic_graphic_coordinates,
graphic_type='new-graphic')
no_coords_layout_graphic = LayoutGraphic(
coordinates=empty_coordinates, graphic_type='empty-coords-graphic')
result = get_layout_document_with_text_and_graphics_replaced_by_graphics(
layout_document,
semantic_graphics=[
SemanticGraphic(layout_graphic=layout_graphic),
SemanticGraphic(layout_graphic=no_coords_layout_graphic)
])
LOGGER.debug('result.pages[0].graphics: %r', result.pages[0].graphics)
assert result.pages[0].graphics[:-1] == [keep_graphic]
LOGGER.debug('result.pages[0].graphics[-1]: %r',
result.pages[0].graphics[-1])
assert result.pages[0].graphics[
-1].graphic_type == layout_graphic.graphic_type
assert result.pages[0].graphics[
-1].coordinates == layout_graphic.coordinates
assert list(
result.pages[0].blocks[0].iter_all_tokens()) == [keep_token]
assert list(
result.pages[0].graphics[-1].related_block.iter_all_tokens()) == [
keep_token, remove_token
]
def get_bounding_box_intersection_area_ratio(
bounding_box_1: BoundingBox,
bounding_box_2: BoundingBox,
empty_ratio: float = 0.0) -> float:
max_area = max(bounding_box_1.area, bounding_box_2.area)
if not max_area:
return empty_ratio
intersection_area = bounding_box_1.intersection(bounding_box_2).area
return intersection_area / max_area
def test_should_find_bbox_and_map_to_page_coordinates( # pylint: disable=too-many-locals
self,
computer_vision_model_mock: MagicMock,
tmp_path: Path,
extract_graphic_assets: bool
):
image_path = tmp_path / 'page10.png'
image = PIL.Image.new('RGB', size=(20, 10), color=(255, 0, 0))
image.save(image_path)
page_images = [DocumentPageImage(
page_number=10,
page_image_path=str(image_path)
)]
layout_document = LayoutDocument(pages=[
_create_page(
coordinates=LayoutPageCoordinates(
x=0, y=0, width=200, height=100, page_number=10
)
)
])
cv_result = computer_vision_model_mock.predict_single.return_value
cv_bbox = BoundingBox(x=1, y=2, width=3, height=4)
cv_result.get_instances_by_type_name.return_value = [
SimpleComputerVisionModelInstance(bounding_box=cv_bbox)
]
expected_page_coordinates = LayoutPageCoordinates(
x=10, y=20, width=30, height=40, page_number=10
)
graphic_provider = ComputerVisionDocumentGraphicProvider(
computer_vision_model=computer_vision_model_mock,
page_image_iterable=page_images,
temp_dir=str(tmp_path)
)
semantic_graphic_list = list(graphic_provider.iter_semantic_graphic_for_layout_document(
layout_document=layout_document,
extract_graphic_assets=extract_graphic_assets
))
assert semantic_graphic_list
semantic_graphic = semantic_graphic_list[0]
LOGGER.debug('semantic_graphic: %s', semantic_graphic)
layout_graphic = semantic_graphic.layout_graphic
assert layout_graphic is not None
assert layout_graphic.coordinates == expected_page_coordinates
if extract_graphic_assets:
assert layout_graphic.local_file_path
assert (
semantic_graphic.relative_path
== os.path.basename(layout_graphic.local_file_path)
)
with PIL.Image.open(layout_graphic.local_file_path) as cropped_image:
assert cropped_image.width == cv_bbox.width
assert cropped_image.height == cv_bbox.height
else:
assert not semantic_graphic.relative_path
def test_should_ignore_graphics_without_coordinates(
self
):
page_graphics = [
LayoutGraphic(coordinates=None)
]
result = get_layout_graphic_with_similar_coordinates(
page_graphics,
BoundingBox(x=10, y=10, width=10, height=1000)
)
assert result is None
def test_should_ignore_matches_below_threshold(
self
):
page_graphics = [
LayoutGraphic(coordinates=LayoutPageCoordinates(
x=10, y=10, width=100, height=100
))
]
result = get_layout_graphic_with_similar_coordinates(
page_graphics,
BoundingBox(x=10, y=10, width=10, height=1000)
)
assert result is None
def is_bounding_box_overlapping_with_any_bounding_boxes(
bounding_box: BoundingBox,
other_bounding_boxes: Sequence[BoundingBox],
max_overlap_ratio: float = 0.1) -> bool:
bounding_box_area = bounding_box.area
for other_bounding_box in other_bounding_boxes:
intersection_bounding_box = bounding_box.intersection(
other_bounding_box)
if not intersection_bounding_box:
continue
if intersection_bounding_box.area / bounding_box_area >= max_overlap_ratio:
return True
return False
def test_should_return_the_best_matching_graphic(
self
):
page_graphics = [
LayoutGraphic(coordinates=LayoutPageCoordinates(
x=10, y=10, width=200, height=100
)),
LayoutGraphic(coordinates=LayoutPageCoordinates(
x=10, y=10, width=100, height=100
)),
LayoutGraphic(coordinates=LayoutPageCoordinates(
x=100, y=10, width=100, height=100
)),
]
result = get_layout_graphic_with_similar_coordinates(
page_graphics,
BoundingBox(x=10, y=10, width=90, height=100)
)
assert result == page_graphics[1]
def test_should_prefer_embedded_graphic( # pylint: disable=too-many-locals
self,
computer_vision_model_mock: MagicMock,
tmp_path: Path
):
image_path = tmp_path / 'page10.png'
image = PIL.Image.new('RGB', size=(20, 10), color=(255, 0, 0))
image.save(image_path)
page_images = [DocumentPageImage(
page_number=10,
page_image_path=str(image_path)
)]
embedded_graphic = LayoutGraphic(
coordinates=LayoutPageCoordinates(
x=10, y=20, width=30, height=40, page_number=10
)
)
layout_document = LayoutDocument(pages=[
_create_page(
coordinates=LayoutPageCoordinates(
x=0, y=0, width=200, height=100, page_number=10
),
graphics=[embedded_graphic]
)
])
cv_result = computer_vision_model_mock.predict_single.return_value
cv_bbox = BoundingBox(x=1, y=2, width=3, height=4)
cv_result.get_instances_by_type_name.return_value = [
SimpleComputerVisionModelInstance(bounding_box=cv_bbox)
]
graphic_provider = ComputerVisionDocumentGraphicProvider(
computer_vision_model=computer_vision_model_mock,
page_image_iterable=page_images,
temp_dir=str(tmp_path)
)
semantic_graphic_list = list(graphic_provider.iter_semantic_graphic_for_layout_document(
layout_document=layout_document,
extract_graphic_assets=True
))
assert semantic_graphic_list
semantic_graphic = semantic_graphic_list[0]
LOGGER.debug('semantic_graphic: %s', semantic_graphic)
assert semantic_graphic.layout_graphic == embedded_graphic
def get_bounding_box_for_layout_parser_coordinates(
coordinates: Tuple[float, float, float, float]) -> BoundingBox:
x1, y1, x2, y2 = coordinates
return BoundingBox(x=x1, y=y1, width=x2 - x1, height=y2 - y1)
def test_should_indicate_not_be_empty_with_non_zero_width_and_height(self):
bounding_box = BoundingBox(0, 0, 100, 100)
assert not bounding_box.is_empty()
assert bounding_box
def test_should_indicate_empty_with_zero_width(self):
bounding_box = BoundingBox(0, 0, 0, 100)
assert bounding_box.is_empty()
assert not bounding_box
def test_should_not_equal_none(self):
assert not BoundingBox(11, 12, 101, 102).__eq__(None)
def test_should_not_equal_bounding_boxes_with_different_height(self):
assert BoundingBox(11, 12, 101, 102) != BoundingBox(11, 12, 101, 999)
def test_should_equal_same_bounding_boxes(self):
assert BoundingBox(11, 12, 101, 102) == BoundingBox(11, 12, 101, 102)
def test_should_calculate_intersection_with_larger_bounding_box(self):
assert (
BoundingBox(110, 120, 50, 60).intersection(
BoundingBox(100, 100, 200, 200)
) == BoundingBox(110, 120, 50, 60)
)
def test_should_calculate_intersection_with_identical_bounding_box(self):
bounding_box = BoundingBox(110, 120, 50, 60)
assert (
bounding_box.intersection(bounding_box) == bounding_box
)
def test_should_scale_by_given_ratio(self):
assert (
BoundingBox(x=1, y=2, width=3, height=4).scale_by(10, 100)
== BoundingBox(x=10, y=200, width=30, height=400)
)
def test_should_calculate_intersection_with_overlapping_bounding_box(self):
assert (
BoundingBox(110, 120, 50, 60).intersection(
BoundingBox(120, 110, 100, 100)
) == BoundingBox(120, 120, 40, 60)
)
LayoutGraphic,
LayoutPage,
LayoutPageCoordinates,
LayoutPageMeta
)
from sciencebeam_parser.cv_models.cv_model import (
SimpleComputerVisionModelInstance
)
from sciencebeam_parser.processors.document_page_image import DocumentPageImage
from sciencebeam_parser.processors.cv_graphic_provider import (
ComputerVisionDocumentGraphicProvider,
get_layout_graphic_with_similar_coordinates
)
BOUNDING_BOX_1 = BoundingBox(x=10, y=10, width=10, height=100)
@pytest.fixture(name='computer_vision_model_mock')
def _computer_vision_model_mock() -> MagicMock:
return MagicMock(name='computer_vision_model')
def _create_page(
coordinates: LayoutPageCoordinates,
graphics: Optional[Sequence[LayoutGraphic]] = None
) -> LayoutPage:
return LayoutPage(
meta=LayoutPageMeta(
page_number=coordinates.page_number,
coordinates=coordinates
def bounding_box(self) -> BoundingBox:
return BoundingBox(x=self.x,
y=self.y,
width=self.width,
height=self.height)
def test_should_calculate_area(self):
bounding_box = BoundingBox(
x=101, y=102, width=200, height=50
)
assert bounding_box.area == 200 * 50
