cv-analysis-service/test/fixtures/figure_detection.py

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