Merge pull request #72 from kotaro-kinoshita/test/add-unittest

add unittest

src/yomitoku/cli/main.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import torch
 from pathlib import Path
 
 import cv2
@@ -235,7 +236,7 @@ def main():
     if args.lite:
         configs["ocr"]["text_recognizer"]["model_name"] = "parseq-small"
 
-        if args.device == "cpu":
+        if args.device == "cpu" or not torch.cuda.is_available():
             configs["ocr"]["text_detector"]["infer_onnx"] = True
 
         # Note: Text Detector以外はONNX推論よりもPyTorch推論の方が速いため、ONNX推論は行わない

src/yomitoku/document_analyzer.py

@@ -127,7 +127,6 @@ def extract_words_within_element(pred_words, element):
     if len(contained_words) == 0:
         return None, None, check_list
 
-    element_direction = "horizontal"
     word_direction = [word.direction for word in contained_words]
     cnt_horizontal = word_direction.count("horizontal")
     cnt_vertical = word_direction.count("vertical")

src/yomitoku/export/export_csv.py

@@ -41,6 +41,8 @@ def save_figure(
     out_path,
     figure_dir="figures",
 ):
+    assert img is not None, "img is required for saving figures"
+
     for i, figure in enumerate(figures):
         x1, y1, x2, y2 = map(int, figure.box)
         figure_img = img[y1:y2, x1:x2, :]

src/yomitoku/export/export_html.py

@@ -110,6 +110,8 @@ def figure_to_html(
     figure_dir="figures",
     width=200,
 ):
+    assert img is not None, "img is required for saving figures"
+
     elements = []
     for i, figure in enumerate(figures):
         x1, y1, x2, y2 = map(int, figure.box)

src/yomitoku/export/export_json.py

@@ -21,6 +21,8 @@ def save_figure(
     out_path,
     figure_dir="figures",
 ):
+    assert img is not None, "img is required for saving figures"
+
     for i, figure in enumerate(figures):
         x1, y1, x2, y2 = map(int, figure.box)
         figure_img = img[y1:y2, x1:x2, :]

src/yomitoku/export/export_markdown.py

@@ -75,6 +75,8 @@ def figure_to_md(
     width=200,
     figure_dir="figures",
 ):
+    assert img is not None, "img is required for saving figures"
+
     elements = []
     for i, figure in enumerate(figures):
         x1, y1, x2, y2 = map(int, figure.box)

tests/test_cli.py

@@ -210,6 +210,7 @@ def test_validate_encoding():
         validate_encoding("utf-9")
 
     assert validate_encoding("utf-8")
+    assert validate_encoding("utf-8-sig")
     assert validate_encoding("shift-jis")
     assert validate_encoding("euc-jp")
     assert validate_encoding("cp932")