[WIP] random plot segments

test/fixtures/page_generation/page.py

@@ -1,5 +1,8 @@
+import io
+import itertools
 import random
 import textwrap
+from itertools import repeat
 from typing import Tuple, Union, Iterable, List
 
 import albumentations as A
@@ -9,6 +12,7 @@ import pytest
 from PIL import Image, ImageOps, ImageFont, ImageDraw
 from PIL.Image import Transpose
 from faker import Faker
+from matplotlib import pyplot as plt
 
 from cv_analysis.utils import star, rconj
 
@@ -69,7 +73,22 @@ Image_t = Union[Image.Image, np.ndarray]
 #     ],
 #     p=0.5,
 # )
-from funcy import juxt, compose, identity, lflatten, lmap, first, iterate, take, last, rest, rcompose
+from funcy import (
+    juxt,
+    compose,
+    identity,
+    lflatten,
+    lmap,
+    first,
+    iterate,
+    take,
+    last,
+    rest,
+    rcompose,
+    pairwise,
+    interleave,
+    keep,
+)
 
 from cv_analysis.locations import TEST_PAGE_TEXTURES_DIR
 
@@ -341,8 +360,59 @@ class ContentGenerator:
     def __init__(self):
         pass
 
-    def __call__(self, boxes: Iterable[Rectangle]) -> Image:
-        return lmap(generate_random_text_block, boxes)
+    def __call__(self, boxes: List[Rectangle]) -> Image:
+        random.shuffle(boxes)
+
+        text_boxes = lmap(generate_random_text_block, every_nth(boxes, 2))
+        plots = lmap(generate_random_plot, every_nth(boxes[1:], 2))
+
+        return text_boxes + plots
+
+
+def every_nth(iterable, n):
+    return itertools.islice(iterable, 0, None, n)
+
+
+def generate_random_plot(rectangle: Rectangle) -> ContentRectangle:
+    block = RandomPlot(*rectangle.coords)
+    block.generate_random_bar_plot(rectangle)
+    return block
+
+
+class RandomPlot(ContentRectangle):
+    def __init__(self, x1, y1, x2, y2):
+        super().__init__(x1, y1, x2, y2)
+
+    def __call__(self, *args, **kwargs):
+        pass
+
+    def generate_random_bar_plot(self, rectangle: Rectangle):
+
+        fig, ax = plt.subplots()
+        x = np.random.randint(low=1, high=11, size=5)
+        y = np.random.randint(low=1, high=11, size=5)
+
+        ax.bar(x, y)
+
+        ax.set_xlabel("X-axis")
+        ax.set_ylabel("Y-axis")
+        ax.set_title("Random Bar Plot")
+        ax.set_facecolor("none")
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format="png", transparent=True)
+        buf.seek(0)
+        image = Image.open(buf)
+        image = image.resize((rectangle.width, rectangle.height))
+        buf.close()
+        plt.close()
+        self.content = image
+
+
+def generate_random_text_block(rectangle: Rectangle) -> ContentRectangle:
+    block = RandomTextBlock(*rectangle.coords)
+    block.generate_random_text(rectangle)
+    return block
 
 
 class RandomTextBlock(ContentRectangle):
@@ -351,6 +421,9 @@ class RandomTextBlock(ContentRectangle):
         self.blank_line_percentage = random.uniform(0, 0.5)
         self.font = ImageFont.load_default()
 
+    def __call__(self, *args, **kwargs):
+        pass
+
     def generate_random_text(self, rectangle: Rectangle):
         def write_line(line, line_number):
             draw.text((0, line_number * text_size), line, font=self.font, fill=(0, 0, 0, 200))
@@ -398,12 +471,6 @@ class RandomTextBlock(ContentRectangle):
         return line
 
 
-def generate_random_text_block(rectangle: Rectangle) -> ContentRectangle:
-    block = RandomTextBlock(*rectangle.coords)
-    block.generate_random_text(rectangle)
-    return block
-
-
 def paste_content(page, content_box: ContentRectangle):
     assert page.mode == "RGB"
     assert content_box.content.mode == "RGBA"