48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
import itertools
|
|
from typing import List
|
|
|
|
from PIL import Image
|
|
from funcy import lsplit, lfilter
|
|
|
|
from cv_analysis.utils import every_nth, zipmap
|
|
from cv_analysis.utils.geometric import is_square_like
|
|
from cv_analysis.utils.merging import merge_related_rectangles
|
|
from cv_analysis.utils.postprocessing import remove_included, remove_overlapping
|
|
from cv_analysis.utils.rectangle import Rectangle
|
|
from synthesis.random import rnd
|
|
from synthesis.segment.segments import (
|
|
generate_random_text_block,
|
|
generate_recursive_random_table_with_caption,
|
|
generate_random_plot_with_caption,
|
|
)
|
|
|
|
|
|
class ContentGenerator:
|
|
def __init__(self):
|
|
self.constrain_layouts = True
|
|
|
|
def __call__(self, boxes: List[Rectangle]) -> Image:
|
|
rnd.shuffle(boxes)
|
|
|
|
figure_boxes, text_boxes = lsplit(is_square_like, boxes)
|
|
|
|
if self.constrain_layouts:
|
|
figure_boxes = merge_related_rectangles(figure_boxes)
|
|
figure_boxes = lfilter(is_square_like, figure_boxes)
|
|
text_boxes = merge_related_rectangles(text_boxes)
|
|
|
|
boxes = list(
|
|
itertools.chain(
|
|
map(generate_random_text_block, every_nth(2, text_boxes)),
|
|
*zipmap(generate_recursive_random_table_with_caption, every_nth(2, text_boxes[1:])),
|
|
*zipmap(generate_recursive_random_table_with_caption, every_nth(2, figure_boxes)),
|
|
*zipmap(generate_random_plot_with_caption, every_nth(2, figure_boxes[1:])),
|
|
)
|
|
)
|
|
|
|
if self.constrain_layouts:
|
|
boxes = remove_included(boxes)
|
|
boxes = remove_overlapping(boxes)
|
|
|
|
return boxes
|