56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import random
|
|
|
|
import pytest
|
|
from funcy import merge
|
|
|
|
from image_prediction.info import Info
|
|
from test.utils.metadata import get_base_position_metadata
|
|
|
|
|
|
@pytest.fixture
|
|
def metadata(images, info_label_map):
|
|
page_idx = 0
|
|
|
|
def current_page_idx():
|
|
nonlocal page_idx
|
|
page_idx += random.randint(0, 3)
|
|
return min(page_idx, len(images) - 1)
|
|
|
|
def build_image_metadata(image):
|
|
width, height = image.size
|
|
page_width = 595
|
|
page_height = 842
|
|
x1 = random.randint(0, page_width - width)
|
|
x2 = x1 + width
|
|
y1 = random.randint(0, page_height - height)
|
|
y2 = y1 + height
|
|
metadata = {
|
|
info_label_map.PAGE_WIDTH: page_width,
|
|
info_label_map.PAGE_HEIGHT: page_height,
|
|
info_label_map.PAGE_IDX: current_page_idx(),
|
|
info_label_map.WIDTH: width,
|
|
info_label_map.HEIGHT: height,
|
|
info_label_map.X1: x1,
|
|
info_label_map.X2: x2,
|
|
info_label_map.Y1: y1,
|
|
info_label_map.Y2: y2,
|
|
info_label_map.ALPHA: image.mode == "RGBA",
|
|
}
|
|
return metadata
|
|
|
|
return list(map(build_image_metadata, images))
|
|
|
|
|
|
@pytest.fixture
|
|
def metadata_formatted(metadata):
|
|
def format_metadata(metadata):
|
|
return {key.value: val for key, val in metadata.items()}
|
|
|
|
return list(map(format_metadata, metadata))
|
|
|
|
|
|
@pytest.fixture
|
|
def base_patch_metadata(width, height, page_width, page_height):
|
|
metadata = get_base_position_metadata(width, height, page_width, page_height)
|
|
metadata = merge(metadata, {Info.X1: 0, Info.Y1: 0, Info.X2: width, Info.Y2: height})
|
|
return metadata |