[WIP] page generation for tests

test/conftest.py

@@ -1,6 +1,7 @@
 pytest_plugins = [
     "test.fixtures.table_parsing",
     "test.fixtures.figure_detection",
+    "test.fixtures.page_generation.page",
 ]
 
 

test/fixtures/page_generation/page.py

@@ -0,0 +1,115 @@
+from typing import Tuple, List
+
+import numpy as np
+import pytest
+from PIL import Image
+
+import cv2 as cv
+
+
+from cv_analysis.utils.display import show_image
+
+Color = Tuple[int, int, int]
+
+
+@pytest.fixture(params=["portrait", "landscape"])
+def orientation(request):
+    return request.param
+
+
+@pytest.fixture(params=[100, 300])
+def dpi(request):
+    return request.param
+
+
+@pytest.fixture(params=["brown", "yellow", "sepia", "gray", "white"])
+def color_name(request):
+    return request.param
+
+
+@pytest.fixture(params=["smooth", "coarse"])
+def texture_name(request):
+    return request.param
+
+
+@pytest.fixture
+def color(color_name):
+    return {
+        "brown": (0.5, 0.3, 0.2),
+        "yellow": (0.5, 0.5, 0.0),
+        "sepia": (0.4, 0.3, 0.2),
+        "gray": (0.3, 0.3, 0.3),
+        "white": (0.0, 0.0, 0.0),
+    }[color_name]
+
+
+@pytest.fixture
+def texture_fn(texture_name):
+    if texture_name == "smooth":
+        return lambda image: cv.blur(image, (5, 5))
+    elif texture_name == "coarse":
+        return lambda image: cv.filter2D(image, -1, np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]]))
+
+
+@pytest.fixture
+def texture(texture_fn, size, color):
+    noise = np.random.rand(*size) * 255
+    noise = color_shift_noise(noise, color)
+    noise = zero_out_below_threshold(noise, 0.1)
+    return Image.fromarray(noise.astype("uint8"))
+
+
+def color_shift_noise(noise: np.ndarray, color: Color):
+    """Creates a 3-tensor from a 2-tensor by stacking the 2-tensor three times weighted by the color tuple."""
+    assert noise.ndim == 2
+    assert isinstance(color, tuple)
+    color = np.array(color)
+    weights = color / color.sum()
+    alpha_channel = np.zeros(noise.shape)
+    colored_noise = np.stack([noise * weight for weight in weights] + [alpha_channel], axis=-1)
+    assert colored_noise.shape == (*noise.shape, 4)
+    return colored_noise
+
+
+@pytest.fixture
+def size(dpi, orientation):
+    if orientation == "portrait":
+        size = (8.5 * dpi, 11 * dpi)
+    else:
+        size = (11 * dpi, 8.5 * dpi)
+    size = tuple(map(int, size))
+    return size
+
+
+@pytest.fixture
+def blank_page(size, texture) -> np.ndarray:
+    """Creates a blank page with a given orientation and dpi."""
+    page = Image.fromarray(np.ones((*size, 3), dtype="uint8") * 255)
+    page = superimpose_texture_with_transparency(page, texture)
+    page = np.array(page)
+    return page
+
+
+def superimpose_texture_with_transparency(page: Image, texture: Image) -> Image:
+    """Superimposes a noise image with transparency onto a page image."""
+    assert page.mode == "RGB"
+    assert texture.mode == "RGBA"
+    assert page.size == texture.size
+    texture = texture.convert("RGBA")
+    page.paste(texture, (0, 0), texture)
+    return page
+
+
+def zero_out_below_threshold(texture, threshold):
+
+    threshold = int(255 * threshold)
+
+    texture[:, :, 0] = (texture[:, :, 0] / texture[:, :, 0].max()) * 255
+    texture[:, :, 1] = (texture[:, :, 1] / texture[:, :, 1].max()) * 255
+    texture[:, :, 2] = (texture[:, :, 2] / texture[:, :, 2].max()) * 255
+    texture[:, :, 3] = 50
+
+    threshold_mask = texture[:, :, 0] >= threshold
+    texture[~threshold_mask] = [0, 0, 0, 0]
+
+    return texture

test/page_generation_test.py

@@ -0,0 +1,6 @@
+from cv_analysis.utils.display import show_image
+
+
+def test_blank_page(blank_page):
+    pass
+    # show_image(blank_page)