fix tests

CDCgov · Aug 30, 2024 · 075b5e0 · 075b5e0
1 parent 4b7806d
commit 075b5e0
Show file tree

Hide file tree

Showing 8 changed files with 816 additions and 77 deletions.
diff --git a/OCR/README.md b/OCR/README.md
@@ -20,7 +20,7 @@ poetry run pytest
 Run benchmark tests
 ```shell
 cd tests
-poetry run pytest bench_test.py -v
+poetry run pytest benchmark_test.py -v
 ```
 
 poetry run pytest bench_test.py -v  

diff --git a/OCR/ocr/main.py b/OCR/ocr/main.py
@@ -23,8 +23,8 @@ def main():
     input_image = args["INPUT-IMAGE-PATH"]
     labels_path = args["LABELS-PATH"]
 
-    segmenter = ImageSegmenter(input_image, segmentation_template, labels_path)
-    segments = segmenter.segment()
+    segmenter = ImageSegmenter()
+    segments = segmenter.load_and_segment(input_image, segmentation_template, labels_path)
 
     print("{:<20} {:<20}".format("Label", "Segment shape"))
     for label, segment in segments.items():

diff --git a/OCR/ocr/services/image_segmenter.py b/OCR/ocr/services/image_segmenter.py
@@ -129,4 +129,4 @@ def load_and_segment(self,
             print(f"raw_image shape: {raw_image.shape}")
             print(f"segmentation_template shape: {segmentation_template.shape}")
 
-        self.segment(raw_image, segmentation_template, labels)
+        return self.segment(raw_image, segmentation_template, labels)
diff --git a/OCR/tests/assets/per_example_marked.pdf b/OCR/tests/assets/per_example_marked.pdf