270 lines
6.5 KiB
Python
270 lines
6.5 KiB
Python
from typing import Tuple
|
|
|
|
import albumentations as A
|
|
import cv2 as cv
|
|
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(
|
|
# [
|
|
# # geometric transforms
|
|
# A.HorizontalFlip(p=0.2),
|
|
# A.RandomRotate90(p=0.2),
|
|
# A.VerticalFlip(p=0.2),
|
|
# # brightness and contrast transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.RandomGamma(p=0.5),
|
|
# A.RandomBrightnessContrast(p=0.5),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# # noise transforms
|
|
# A.SomeOf(
|
|
# [
|
|
# A.Emboss(p=0.05),
|
|
# A.ImageCompression(p=0.05),
|
|
# A.PixelDropout(p=0.05),
|
|
# ],
|
|
# p=0.5,
|
|
# n=2,
|
|
# ),
|
|
# # color transforms
|
|
# A.SomeOf(
|
|
# [
|
|
# A.ColorJitter(p=1),
|
|
# A.RGBShift(p=1, r_shift_limit=0.1, g_shift_limit=0.1, b_shift_limit=0.1),
|
|
# A.ChannelShuffle(p=1),
|
|
# ],
|
|
# p=0.5,
|
|
# 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.5,
|
|
# ),
|
|
# # environmental transforms
|
|
# A.OneOf(
|
|
# [
|
|
# A.RandomRain(p=0.2, rain_type="drizzle"),
|
|
# A.RandomFog(p=0.2, fog_coef_upper=0.4),
|
|
# A.RandomSnow(p=0.2),
|
|
# ],
|
|
# p=0.5,
|
|
# ),
|
|
# ],
|
|
# p=0.5,
|
|
# )
|
|
|
|
transform = A.Compose(
|
|
[
|
|
# A.ColorJitter(p=1),
|
|
]
|
|
)
|
|
|
|
|
|
Color = Tuple[int, int, int]
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
"portrait",
|
|
# "landscape",
|
|
]
|
|
)
|
|
def orientation(request):
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
# 30,
|
|
100,
|
|
]
|
|
)
|
|
def dpi(request):
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
"brown",
|
|
# "yellow",
|
|
# "sepia",
|
|
# "gray",
|
|
# "white",
|
|
# "light_red",
|
|
# "light_blue",
|
|
]
|
|
)
|
|
def color_name(request):
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
"smooth",
|
|
"coarse",
|
|
"neutral",
|
|
]
|
|
)
|
|
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
|
|
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
|
|
|
|
|
|
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 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
|
|
def size(dpi, orientation):
|
|
if orientation == "portrait":
|
|
size = (8.5 * dpi, 11 * dpi)
|
|
elif orientation == "landscape":
|
|
size = (11 * dpi, 8.5 * dpi)
|
|
else:
|
|
raise ValueError(f"Unknown orientation: {orientation}")
|
|
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.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 texture.mode == "RGBA"
|
|
assert page.size == texture.size
|
|
page.paste(texture, (0, 0), texture)
|
|
return page
|