🎨 add list_image_files and assert_len

dobble/card.py

@@ -13,6 +13,8 @@
 from tqdm import tqdm
 
 from dobble.optim import get_cards
+from dobble.utils import assert_len
+from dobble.utils import list_image_files
 from dobble.utils import new_folder
 
 DEBUG = False
@@ -237,11 +239,11 @@ def main(masks_folder: str,
         circle_width_pix: Width of the circle around each card. Covariant with card_size_pix
         n_iter: Number of evolution steps for each card
     """
-    names = [f.name for f in os.scandir(masks_folder)]
-    assert len(names) == 57
+    names = list_image_files(masks_folder)
+    assert_len(names, 57)
 
     cards = get_cards()
-    assert len(cards) == 57
+    assert_len(cards, 57)
 
     new_folder(out_cards_folder)
 

dobble/mask.py

@@ -10,6 +10,7 @@
 from tqdm import tqdm
 
 from dobble.utils import new_folder
+from dobble.utils import list_image_files
 
 DEBUG = False
 
@@ -59,7 +60,7 @@ def main(symbols_folder: str,
         margin_pix: Dilation applied around the mask, covariant with computing_size_pix
         ths: Pixels the intensity of which is above this threshold are considered as white background
     """
-    names = [f.name for f in os.scandir(symbols_folder)]
+    names = list_image_files(symbols_folder)
 
     new_folder(out_masks_folder)
 

dobble/pdf.py

@@ -11,6 +11,8 @@
 from tqdm import tqdm
 
 from dobble.utils import new_folder
+from dobble.utils import assert_len
+from dobble.utils import list_image_files
 
 
 def main(cards_folder: str,
@@ -24,8 +26,8 @@ def main(cards_folder: str,
         out_print_folder: Output folder containing the batched cards and the PDF file
         card_size_cm: Diameter of the output Dobble cards to print
     """
-    names = [f.name for f in os.scandir(cards_folder)]
-    assert len(names) == 57
+    names = list_image_files(cards_folder)
+    assert_len(names, 57)
     names += [None]  # Pad to have an even size
 
     pdf_path = os.path.join(out_print_folder, "cards.pdf")

dobble/preprocess.py

@@ -2,6 +2,7 @@
 """Make all images square and rotation-proof"""
 
 
+from dobble.utils import assert_len, list_image_files
 import matplotlib.pyplot as plt
 import math
 import os
@@ -52,14 +53,6 @@ def _set_white_background(img: np.ndarray) -> np.ndarray:
     return img
 
 
-def list_image_files(images_folder: str) -> List[str]:
-    """List image files."""
-    image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff')
-    return [f.name
-            for f in os.scandir(images_folder)
-            if f.name.lower().endswith(image_extensions)]
-
-
 def main(images_folder: str,
          out_images_folder: str):
     """

dobble/utils.py

@@ -4,9 +4,24 @@
 
 import os
 import shutil
+from typing import Sized, List
 
 
 def new_folder(folder: str):
     if os.path.exists(folder):
         shutil.rmtree(folder, ignore_errors=True)
     os.makedirs(folder)
+
+
+def assert_len(seq: Sized, size: int):
+    """Assert Python list has expected length."""
+    assert len(seq) == size, \
+        f"Expect sequence of length {size}. Got length {len(seq)}."
+
+
+def list_image_files(images_folder: str) -> List[str]:
+    """List image files."""
+    image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff')
+    return [f.name
+            for f in os.scandir(images_folder)
+            if f.name.lower().endswith(image_extensions)]