Add binarization to test_chat_models.py

colibrisson · Sep 7, 2023 · 0857e02 · 0857e02
1 parent ceb1a34
commit 0857e02
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/test/test_chat_models.py b/test/test_chat_models.py
@@ -58,6 +58,10 @@ def check_line_direction(baseline_seg: dict) -> dict:
         # Load image
         img = Image.open(img_path)
 
+        # Binarize the image
+        img = img.convert("L")
+        img = img.point(lambda x: 0 if x < 128 else 255, "1")
+
         # Segment the image
         baseline_seg = blla.segment(img, text_direction="vertical-rl", model=seg_model)