123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
import abc
|
|
import random
|
|
from copy import deepcopy
|
|
from itertools import chain
|
|
from operator import itemgetter
|
|
|
|
import fpdf
|
|
import pytest
|
|
from funcy import juxt, merge, rpartial, compose
|
|
|
|
from image_prediction.image_extractor.extractor import ImageMetadataPair
|
|
from image_prediction.info import Info
|
|
from test.conftest import get_base_position_metadata, add_image, random_single_color_image_from_metadata
|
|
|
|
|
|
def test_image_stitcher(partial_image_metadata_pairs):
|
|
pass
|
|
|
|
|
|
@pytest.mark.parametrize("width", [160])
|
|
@pytest.mark.parametrize("height", [90])
|
|
@pytest.mark.parametrize("page_width", [int(160 * 1.1)])
|
|
@pytest.mark.parametrize("page_height", [int(90 * 1.1)])
|
|
def test_partial_image_metadata_pairs(patches_metadata, page_width, page_height):
|
|
pdf = fpdf.FPDF(unit="pt", format=(page_width, page_height))
|
|
|
|
for patch in patches_metadata:
|
|
image = random_single_color_image_from_metadata(patch)
|
|
add_image(pdf, ImageMetadataPair(image, patch))
|
|
|
|
pdf.output("/tmp/bla.pdf")
|
|
|
|
|
|
class SplitKeyMapper(abc.ABC):
|
|
def __init__(self, box: dict, keymap: dict):
|
|
self.box = box
|
|
self.keymap = keymap
|
|
|
|
def __getitem__(self, item):
|
|
return self.box[self.keymap[item]]
|
|
|
|
def __setitem__(self, key, value):
|
|
self.box[self.keymap[key]] = value
|
|
|
|
|
|
class HorizontalKeyMapper(SplitKeyMapper):
|
|
def __init__(self, box: dict):
|
|
super().__init__(box, {"dim": Info.WIDTH, "c1": Info.X1, "c2": Info.X2})
|
|
|
|
|
|
class VerticalKeyMapper(SplitKeyMapper):
|
|
def __init__(self, box: dict):
|
|
super().__init__(box, {"dim": Info.HEIGHT, "c1": Info.Y1, "c2": Info.Y2})
|
|
|
|
|
|
class BoxSplitter:
|
|
def __init__(self):
|
|
self.__steps = None
|
|
|
|
def split_box(self, box, steps=5):
|
|
self.__steps = steps
|
|
return self.__split_recursively(box, 0)
|
|
|
|
def __split_recursively(self, box, step):
|
|
return self.__split_and_recurse(box, step) if self.__steps_left(step) else self.__base_case(box)
|
|
|
|
def __steps_left(self, step):
|
|
return step < self.__steps
|
|
|
|
@staticmethod
|
|
def __base_case(box):
|
|
return [box]
|
|
|
|
def __split_and_recurse(self, box, step):
|
|
new_boxes = self.__random_split(box)
|
|
new_boxes_per_branch = self.__tree_recurse(new_boxes, step + 1)
|
|
return chain.from_iterable(new_boxes_per_branch)
|
|
|
|
def __random_split(self, box):
|
|
splitter = random.choice([self.__split_horizontal, self.__split_vertical])
|
|
new_boxes = splitter(box)
|
|
return new_boxes
|
|
|
|
def __tree_recurse(self, boxes, step):
|
|
return map(rpartial(self.__split_recursively, step + 1), boxes)
|
|
|
|
def __split_horizontal(self, box):
|
|
return self.__split_if_large_enough(HorizontalKeyMapper(box))
|
|
|
|
def __split_vertical(self, box):
|
|
return self.__split_if_large_enough(VerticalKeyMapper(box))
|
|
|
|
def __split_if_large_enough(self, sabox: SplitKeyMapper):
|
|
return self.__get_child_boxes(sabox) if self.__large_enough(sabox) else self.__base_case(sabox.box)
|
|
|
|
@staticmethod
|
|
def __large_enough(box):
|
|
return box["dim"] >= 10
|
|
|
|
@staticmethod
|
|
def __get_child_boxes(sabox: SplitKeyMapper):
|
|
|
|
split_len = random.randint(5, sabox["dim"] - 5)
|
|
split_point = sabox["c1"] + split_len
|
|
|
|
box_left, box_right = juxt(deepcopy, deepcopy)(sabox)
|
|
|
|
box_left["dim"] = split_len
|
|
box_right["dim"] = sabox["dim"] - split_len
|
|
|
|
box_left["c2"] = split_point
|
|
box_right["c1"] = split_point
|
|
|
|
return box_left.box, box_right.box
|
|
|
|
|
|
@pytest.fixture()
|
|
def patches_metadata(width, height, page_width, page_height):
|
|
box = get_base_position_metadata(width, height, page_width, page_height)
|
|
box = merge(box, {Info.X1: 0, Info.Y1: 0, Info.X2: width, Info.Y2: height})
|
|
boxes = BoxSplitter().split_box(box)
|
|
return boxes
|