import textwrap import cv2 import numpy as np import pytest from lorem_text import lorem from cv_analysis.figure_detection.figure_detection import ( detect_figures, ) @pytest.fixture def page_with_images(random_image, background): page_image = paste_image(background, random_image, (200, 200)) 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 detect_figures 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): h, w = image.shape[:2] x, y = coords page_image[x : x + h, y : y + w] = image return page_image