diff --git a/cv_analysis/locations.py b/cv_analysis/locations.py index 4bf06bd..27dc749 100644 --- a/cv_analysis/locations.py +++ b/cv_analysis/locations.py @@ -5,5 +5,8 @@ from pathlib import Path MODULE_PATH = Path(__file__).resolve().parents[0] PACKAGE_ROOT_PATH = MODULE_PATH.parents[0] REPO_ROOT_PATH = PACKAGE_ROOT_PATH + TEST_DIR_PATH = REPO_ROOT_PATH / "test" -TEST_DATA_DVC = TEST_DIR_PATH / "test_data.dvc" +TEST_DATA_DVC = TEST_DIR_PATH / "test_data.dvc" # TODO: remove once new tests are in place +TEST_DATA_DIR = TEST_DIR_PATH / "data" +TEST_PAGE_TEXTURES_DIR = TEST_DATA_DIR / "paper" diff --git a/test/fixtures/page_generation/page.py b/test/fixtures/page_generation/page.py index 7d0dcc4..d97cf31 100644 --- a/test/fixtures/page_generation/page.py +++ b/test/fixtures/page_generation/page.py @@ -1,10 +1,10 @@ +import random from typing import Tuple import albumentations as A -import cv2 as cv import numpy as np import pytest -from PIL import Image +from PIL import Image, ImageOps # transform = A.Compose( # [ @@ -37,7 +37,8 @@ from PIL import Image # # ), # ] # ) -from funcy import compose, identity, juxt + +from PIL.Image import Transpose # # transform = A.Compose( @@ -95,6 +96,7 @@ from funcy import compose, identity, juxt # ], # p=0.5, # ) +from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR transform = A.Compose( [ @@ -106,10 +108,24 @@ transform = A.Compose( Color = Tuple[int, int, int] +@pytest.fixture( + params=[ + # "rough_grain", + # "plain", + # "digital", + "crumpled", + ] +) +def base_texture(request, size): + texture = Image.open(TEST_PAGE_TEXTURES_DIR / (request.param + ".jpg")) + texture = texture.resize(size) + return texture + + @pytest.fixture( params=[ "portrait", - # "landscape", + "landscape", ] ) def orientation(request): @@ -128,9 +144,8 @@ def dpi(request): @pytest.fixture( params=[ - "brown", - # "yellow", - # "sepia", + # "brown", + "sepia", # "gray", # "white", # "light_red", @@ -152,91 +167,161 @@ 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": (173, 155, 109), - "gray": (0.3, 0.3, 0.3), - "white": (0.0, 0.0, 0.0), - "light_red": (0.5, 0.0, 0.0), - "light_blue": (0.0, 0.0, 0.5), - }[color_name] +@pytest.fixture( + params=[ + 30, + 70, + 150, + ] +) +def color_intensity(request): + return request.param -@pytest.fixture -def texture_fn(texture_name, size): - if texture_name == "smooth": - fn = blur - elif texture_name == "coarse": - 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) +def random_flip(image): + if random.choice([True, False]): + image = image.transpose(Transpose.FLIP_LEFT_RIGHT) + if random.choice([True, False]): + image = image.transpose(Transpose.FLIP_TOP_BOTTOM) return image +# @pytest.fixture +# def color(color_name): +# return { +# "brown": (0.5, 0.3, 0.2), +# "yellow": (0.5, 0.5, 0.0), +# "sepia": (173, 155, 109), +# "gray": (0.3, 0.3, 0.3), +# "white": (0.0, 0.0, 0.0), +# "light_red": (0.5, 0.0, 0.0), +# "light_blue": (0.0, 0.0, 0.5), +# }[color_name] + + @pytest.fixture -def texture(texture_fn, size, color): - noise_arr = np.random.rand(*size) * 255 - noise_arr = color_shift_noise(noise_arr, color) - - noise_arr = zero_out_below_threshold(noise_arr, 0.4) - noise_arr = texture_fn(noise_arr) - assert noise_arr.max() <= 255 - - noise_img = Image.fromarray(noise_arr) - # noinspection PyTypeChecker - assert np.equal(noise_arr, np.array(noise_img)).all() - return noise_img +def color(color_name): + return { + "brown": "#7d6c5b", + "sepia": "#b8af88", + "gray": "#9c9c9c", + "white": "#ffffff", + "light_red": "#d68c8b", + "light_blue": "#8bd6d6", + }[color_name] -def color_shift_noise(noise: np.ndarray, color: Color): +# @pytest.fixture +# def texture_fn(texture_name, size): +# if texture_name == "smooth": +# fn = blur +# elif texture_name == "coarse": +# 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 = color_shift_noise(noise_arr, color) +# +# noise_arr = zero_out_below_threshold(noise_arr, 0.4) +# noise_arr = texture_fn(noise_arr) +# assert noise_arr.max() <= 255 +# +# noise_img = Image.fromarray(noise_arr) +# # noinspection PyTypeChecker +# assert np.equal(noise_arr, np.array(noise_img)).all() +# return noise_img + + +@pytest.fixture +def texture(base_texture, color, color_intensity): + color_image = Image.new("RGBA", base_texture.size, color) + color_image.putalpha(color_intensity) + texture = superimpose_texture_with_transparency(base_texture, color_image) + return texture + + +def tint_image(src, color="#FFFFFF"): + src.load() + r, g, b, alpha = src.split() + gray = ImageOps.grayscale(src) + result = ImageOps.colorize(gray, (0, 0, 0), color) + result.putalpha(alpha) + return result + + +# 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) +# assert max(color) <= 255 +# assert noise.max() <= 255 +# +# color = np.array(color) +# weights = color / color.sum() +# assert max(weights) <= 1 +# +# alpha_channel = np.ones(noise.shape) * 255 +# colored_noise = np.stack([noise * weight for weight in weights] + [alpha_channel], axis=-1).astype(np.uint8) +# +# assert colored_noise.shape == (*noise.shape, 4) +# +# return colored_noise + + +def color_shift_array(image: 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 image.ndim == 3 + assert image.shape[-1] == 3 assert isinstance(color, tuple) assert max(color) <= 255 - assert noise.max() <= 255 + assert image.max() <= 255 color = np.array(color) - weights = color / color.sum() + weights = color / color.sum() / 10 assert max(weights) <= 1 - alpha_channel = np.ones(noise.shape) * 255 - colored_noise = np.stack([noise * weight for weight in weights] + [alpha_channel], axis=-1).astype(np.uint8) + # alpha_channel = np.ones(image.shape) * 255 + # colored_noise = np.stack([image * weight for weight in weights] + [alpha_channel], axis=-1).astype(np.uint8) + # colored = np.column_stack([image * weights] + [alpha_channel]).astype(np.uint8) + colored = (image * weights).astype(np.uint8) - assert colored_noise.shape == (*noise.shape, 4) + assert colored.shape == image.shape - return colored_noise + return colored -def zero_out_below_threshold(texture, threshold): - assert texture.max() <= 255 - - texture[:, :, 3] = 100 - - threshold = int(texture[:, :, 0:3].sum(axis=2).max() * threshold) - threshold_mask = texture[:, :, 0:3].sum(axis=2) >= threshold - texture[~threshold_mask] = [0, 0, 0, 50] - - return texture +# def zero_out_below_threshold(texture, threshold): +# assert texture.max() <= 255 +# +# texture[:, :, 3] = 100 +# +# threshold = int(texture[:, :, 0:3].sum(axis=2).max() * threshold) +# threshold_mask = texture[:, :, 0:3].sum(axis=2) >= threshold +# texture[~threshold_mask] = [0, 0, 0, 50] +# +# return texture @pytest.fixture @@ -251,19 +336,26 @@ def size(dpi, orientation): 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.zeros((*size, 4), dtype=np.uint8) * 255) - page = superimpose_texture_with_transparency(page, texture) - # page = transform(image=page)["image"] - return page - +# @pytest.fixture +# 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) +# # page = transform(image=page)["image"] +# 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 == "RGBA" + assert page.mode == "RGB" assert texture.mode == "RGBA" assert page.size == texture.size page.paste(texture, (0, 0), texture) return page + + +@pytest.fixture +def blank_page(texture) -> np.ndarray: + """Creates a blank page with a given orientation and dpi.""" + page = random_flip(texture) + return page