diff --git a/ebl_ai/generate_bounding_boxes_resource.py b/ebl_ai/generate_bounding_boxes_resource.py index 9c48809..8b4f5d3 100644 --- a/ebl_ai/generate_bounding_boxes_resource.py +++ b/ebl_ai/generate_bounding_boxes_resource.py @@ -4,6 +4,7 @@ import falcon import numpy as np from PIL import Image + Image.MAX_IMAGE_PIXELS = None from ebl_ai.model import Model diff --git a/ebl_ai/model.py b/ebl_ai/model.py index e234f99..8cde5bd 100644 --- a/ebl_ai/model.py +++ b/ebl_ai/model.py @@ -23,7 +23,6 @@ class Model: def __init__(self, configFile: str, checkpoint: str): self.model = MMOCRInferencer(det=configFile, det_weights=checkpoint) - def _predict(self, image_path: str) -> List[List[float]]: x = self.model(image_path)["predictions"] result = [] @@ -72,4 +71,3 @@ def predict( boundary_results = self._predict(file.name) return self._polygons_with_probabilites_to_rectangle(boundary_results) - diff --git a/ebl_ai/tests/test_model.py b/ebl_ai/tests/test_model.py index 195c3d3..83e47ac 100644 --- a/ebl_ai/tests/test_model.py +++ b/ebl_ai/tests/test_model.py @@ -12,7 +12,6 @@ TEST_IMAGE_PATH = "ebl_ai/tests/test_image.jpg" -#@pytest.mark.skip(reason="Takes very long") def test_model_predictions(): model = Model(configFile=CONFIG_FILE, checkpoint=CHECKPOINT) @@ -21,9 +20,7 @@ def test_model_predictions(): assert isinstance(predictions[0], BoundingBoxesPrediction) assert len(predictions) > 1 - model.show_result(TEST_IMAGE_PATH, "./test_image_prediction.jpg", False) - img = Image.open(TEST_IMAGE_PATH) - predictions = model.predict(np.asarray(img)) + predictions = model.model(np.asarray(img), show=True) assert isinstance(predictions[0], BoundingBoxesPrediction) assert len(predictions) > 1