cv-analysis-service/test/fixtures/figure_detection.py
Julius Unverfehrt ce9e92876c Pull request #16: Add table parsing fixtures
Merge in RR/cv-analysis from add_table_parsing_fixtures to master

Squashed commit of the following:

commit cfc89b421b61082c8e92e1971c9d0bf4490fa07e
Merge: a7ecb05 73c66a8
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Mon Jul 11 12:19:01 2022 +0200

    Merge branch 'master' of ssh://git.iqser.com:2222/rr/cv-analysis into add_table_parsing_fixtures

commit a7ecb05b7d8327f0c7429180f63a380b61b06bc3
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Mon Jul 11 12:02:07 2022 +0200

    refactor

commit 466f217e5a9ee5c54fd38c6acd28d54fc38ff9bb
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Mon Jul 11 10:24:14 2022 +0200

    deleted unused imports and unused lines of code

commit c58955c8658d0631cdd1c24c8556d399e3fd9990
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Mon Jul 11 10:16:01 2022 +0200

    black reformatted files

commit f8bcb10a00ff7f0da49b80c1609b17997411985a
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Tue Jul 5 15:15:00 2022 +0200

    reformat files

commit 432e8a569fd70bd0745ce0549c2bfd2f2e907763
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Tue Jul 5 15:08:22 2022 +0200

    added better test for generic pages with table WIP as thicker lines create inconsistent results.
    added test for patchy tables which does not work yet

commit 2aac9ebf5c76bd963f8c136fe5dd4c2d7681b469
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Mon Jul 4 16:56:29 2022 +0200

    added new fixtures for table_parsing_test.py

commit 37606cac0301b13e99be2c16d95867477f29e7c4
Author: llocarnini <lillian.locarnini@iqser.com>
Date:   Fri Jul 1 16:02:44 2022 +0200

    added separate file for table parsing fixtures, where fixtures for generic tables were added. WIP tests for generic table fixtures
2022-07-11 12:25:16 +02:00

101 lines
3.0 KiB
Python

import textwrap
import cv2
import numpy as np
import pytest
from PIL import Image
from lorem_text import lorem
from funcy import first
from cv_analysis.figure_detection.figure_detection_pipeline import (
make_figure_detection_pipeline,
)
from cv_analysis.utils.display import show_mpl
@pytest.fixture
def page_with_images(random_image, n_images, background):
page_image = Image.fromarray(background.astype("uint8")).convert("RGB")
page_image = paste_image(page_image, random_image, (200, 200))
if n_images == 2:
page_image = paste_image(page_image, random_image, (1000, 2600))
return np.array(page_image)
@pytest.fixture
def page_with_text(background, font_scale, font_style, text_types):
cursor = (50, 50)
image = background
body_height = image.shape[0] // 3
if "header" in text_types:
image = paste_text(image, cursor, font_scale, font_style, y_stop=70)
if "body" in text_types:
cursor = (image.shape[1] // 2, 70)
image = paste_text(image, cursor, font_scale, font_style, y_stop=body_height)
cursor = (50, body_height + 70)
image = paste_text(
image, cursor, font_scale, font_style, y_stop=body_height * 2
)
if "caption" in text_types:
cursor = (image.shape[1] // 2, image.shape[0] - 100)
image = paste_text(
image, cursor, font_scale, font_style, y_stop=body_height * 3
)
return image
@pytest.fixture
def page_with_images_and_text(page_with_images, page_with_text):
return np.fmin(page_with_text, page_with_images)
@pytest.fixture
def background(background_color):
return np.ones((3508, 2480, 3), dtype="uint8") * background_color
@pytest.fixture
def random_image(image_size):
return np.random.rand(*image_size, 3) * 255
@pytest.fixture
def figure_detection_pipeline():
return make_figure_detection_pipeline()
def paste_text(image: np.ndarray, cursor, font_scale, font_style, y_stop):
def paste_text_at_cursor(x_start, y_start, y_stop):
# TODO: adjust incorrect right margin
text = lorem.paragraphs(1) * 200
(dx, dy), base = cv2.getTextSize(
text, fontFace=font_style, fontScale=font_scale, thickness=1
)
dy += base
# char_width = dx // len(text)
text = textwrap.fill(text=text, width=(dx // page_width))
for i, line in enumerate(text.split("\n")):
y = y_start + i * dy
if y > y_stop:
break
cv2.putText(
image,
line,
org=(x_start, y),
fontFace=font_style,
fontScale=font_scale,
color=(0, 0, 0),
thickness=1,
lineType=cv2.LINE_AA,
)
x_start, y_start = cursor
page_width = image.shape[1]
paste_text_at_cursor(x_start, y_start, y_stop)
return image
def paste_image(page_image, image, coords):
image = Image.fromarray(image.astype("uint8")).convert("RGBA")
page_image.paste(image, coords)
return page_image