Improve page texture logic

test/fixtures/page_generation/page.py

@@ -6,6 +6,39 @@ import numpy as np
 import pytest
 from PIL import Image
 
+# transform = A.Compose(
+#     [
+#         # brightness and contrast transforms
+#         A.OneOf(
+#             [
+#                 A.RandomGamma(p=0.2),
+#                 A.RandomBrightnessContrast(p=0.2, brightness_limit=0.05, contrast_limit=0.05),
+#             ],
+#             p=0.5,
+#         ),
+#         # color transforms
+#         A.SomeOf(
+#             [
+#                 A.ColorJitter(p=1),
+#                 A.RGBShift(p=1, r_shift_limit=0.3, g_shift_limit=0.3, b_shift_limit=0.3),
+#                 A.ChannelShuffle(p=1),
+#             ],
+#             p=1.0,
+#             n=3,  # 3 => all
+#         ),
+#         # # blurring and sharpening transforms
+#         # A.OneOf(
+#         #     [
+#         #         A.GaussianBlur(p=0.05),
+#         #         A.MotionBlur(p=0.05, blur_limit=21),
+#         #         A.Sharpen(p=0.05),
+#         #     ],
+#         #     p=0.0,
+#         # ),
+#     ]
+# )
+from funcy import compose, identity, juxt
+
 #
 # transform = A.Compose(
 #     [
@@ -63,40 +96,6 @@ from PIL import Image
 #     p=0.5,
 # )
 
-
-# transform = A.Compose(
-#     [
-#         # brightness and contrast transforms
-#         A.OneOf(
-#             [
-#                 A.RandomGamma(p=0.2),
-#                 A.RandomBrightnessContrast(p=0.2, brightness_limit=0.05, contrast_limit=0.05),
-#             ],
-#             p=0.5,
-#         ),
-#         # color transforms
-#         A.SomeOf(
-#             [
-#                 A.ColorJitter(p=1),
-#                 A.RGBShift(p=1, r_shift_limit=0.3, g_shift_limit=0.3, b_shift_limit=0.3),
-#                 A.ChannelShuffle(p=1),
-#             ],
-#             p=1.0,
-#             n=3,  # 3 => all
-#         ),
-#         # # blurring and sharpening transforms
-#         # A.OneOf(
-#         #     [
-#         #         A.GaussianBlur(p=0.05),
-#         #         A.MotionBlur(p=0.05, blur_limit=21),
-#         #         A.Sharpen(p=0.05),
-#         #     ],
-#         #     p=0.0,
-#         # ),
-#     ]
-# )
-from cv_analysis.utils.display import show_image
-
 transform = A.Compose(
     [
         # A.ColorJitter(p=1),
@@ -144,8 +143,8 @@ def color_name(request):
 
 @pytest.fixture(
     params=[
-        # "smooth",
-        # "coarse",
+        "smooth",
+        "coarse",
         "neutral",
     ]
 )
@@ -169,26 +168,41 @@ def color(color_name):
 @pytest.fixture
 def texture_fn(texture_name, size):
     if texture_name == "smooth":
-        return lambda image: cv.blur(image, np.array(size) // 100)
+        fn = blur
     elif texture_name == "coarse":
-        return lambda image: cv.filter2D(image, -1, np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]]))
-    elif texture_name == "neutral":
-        return lambda image: image
+        fn = compose(overlay, juxt(blur, sharpen))
+    else:
+        fn = identity
+
+    return fn
+
+
+def blur(image: np.ndarray):
+    return cv.blur(image, (3, 3))
+
+
+def sharpen(image: np.ndarray):
+    return cv.filter2D(image, -1, np.array([[-1, -1, -1], [-1, 6, -1], [-1, -1, -1]]))
+
+
+def overlay(images, mode=np.sum):
+    assert mode in [np.sum, np.max]
+    images = np.stack(list(images))
+    image = mode(images, axis=0)
+    image = (image / image.max() * 255).astype(np.uint8)
+    return image
 
 
 @pytest.fixture
 def texture(texture_fn, size, color):
     noise_arr = np.random.rand(*size) * 255
-    noise_arr = texture_fn(noise_arr)
     noise_arr = color_shift_noise(noise_arr, color)
 
-    show_image(noise_arr, backend="mpl", title="noise1")
     noise_arr = zero_out_below_threshold(noise_arr, 0.4)
+    noise_arr = texture_fn(noise_arr)
     assert noise_arr.max() <= 255
 
-    show_image(noise_arr, backend="mpl", title="noise2")
     noise_img = Image.fromarray(noise_arr)
-    show_image(noise_img, backend="mpl", title="noise3")
     # noinspection PyTypeChecker
     assert np.equal(noise_arr, np.array(noise_img)).all()
     return noise_img
@@ -216,12 +230,11 @@ def color_shift_noise(noise: np.ndarray, color: Color):
 def zero_out_below_threshold(texture, threshold):
     assert texture.max() <= 255
 
+    texture[:, :, 3] = 100
+
     threshold = int(texture[:, :, 0:3].sum(axis=2).max() * threshold)
-
-    # texture[:, :, 3] = 100
-
     threshold_mask = texture[:, :, 0:3].sum(axis=2) >= threshold
-    texture[~threshold_mask] = [0, 0, 0, 0]
+    texture[~threshold_mask] = [0, 0, 0, 50]
 
     return texture
 
@@ -243,7 +256,6 @@ def blank_page(size, texture) -> np.ndarray:
     """Creates a blank page with a given orientation and dpi."""
     page = Image.fromarray(np.zeros((*size, 4), dtype=np.uint8) * 255)
     page = superimpose_texture_with_transparency(page, texture)
-    show_image(page, backend="mpl", title="page")
     # page = transform(image=page)["image"]
     return page
 
@@ -253,6 +265,5 @@ def superimpose_texture_with_transparency(page: Image, texture: Image) -> Image:
     assert page.mode == "RGBA"
     assert texture.mode == "RGBA"
     assert page.size == texture.size
-    show_image(texture, backend="mpl", title="texture")
     page.paste(texture, (0, 0), texture)
     return page