Add row generation

comic_layout/comic_layout.py

@@ -5,29 +5,42 @@ from utils import profile
 
 
 class LayoutGenerator():
+
+    @classmethod
+    def _gen_image_row(cls, images, width):
+        image_row = np.hstack(images)
+        if image_row.shape[1] != width:
+            nominal_height = image_row.shape[0]
+            image_row = cv2.resize(
+                image_row,
+                (width, (nominal_height * width) // nominal_height)
+            )
+        return image_row
+
+    @classmethod
+    def _get_row_layout(cls, scores):
+        return [2, 3, 1, 3, 1]
+
     @classmethod
     @profile
-    def get_layout(cls, frames):
-        result_imgs = cls._pad_images(frames)
+    def get_layout(cls, frames, scores):
+        result_images = cls._pad_images(frames)
+        row_layout = cls._get_row_layout(scores)
+        width = result_images[0].shape[1]
 
-        first_row = np.hstack(result_imgs[:2])
-        second_row = np.hstack(result_imgs[2:5])
-        third_row = np.hstack(result_imgs[5:7])
-        fourth_row = np.hstack(result_imgs[7:10])
+        index = 0
+        rows = []
+        for row_length in row_layout:
+            row = cls._gen_image_row(result_images[index:index + row_length], width)
+            rows.append(row)
+            index += row_length
 
-        second_row = cv2.resize(second_row,
-                                (first_row.shape[1],
-                                 (second_row.shape[0] * first_row.shape[1]) // second_row.shape[1]))
-        fourth_row = cv2.resize(fourth_row,
-                                (first_row.shape[1],
-                                 (fourth_row.shape[0] * first_row.shape[1]) // fourth_row.shape[1]))
-
-        return np.vstack([first_row, second_row, third_row, fourth_row])
+        return np.vstack(rows)
 
     @staticmethod
     def _pad_images(frames):
-        padded_result_imgs = []
+        padded_result_images = []
         for img in frames:
             padded_img = cv2.copyMakeBorder(img, 5, 5, 5, 5, cv2.BORDER_CONSTANT, value=(255, 255, 255))
-            padded_result_imgs.append(padded_img)
-        return padded_result_imgs
+            padded_result_images.append(padded_img)
+        return padded_result_images