diff --git a/README.md b/README.md index 42e37351d3fb53ebbd1c76cbc7e55f3bf3563ffd..5d50b36c932f0da4dd18fe47cb4a736055149266 100644 --- a/README.md +++ b/README.md @@ -63,9 +63,9 @@ image = cv2.cvtColor(cv2.imread(IMAGE_PATH), cv2.COLOR_BGR2RGB) Then one can initialize and load the trained model with the parameters used during training. The directory passed as parameter should have: -- a single model file `*.pt`, -- a single charset file `*.pkl`, -- a single parameters file `*parameters.yml`. +- a `model.pt` file, +- a `charset.pkl` file, +- a `parameters.yml` file corresponding to the `inference_parameters.yml` file generated during training. ```python model_path = "models" diff --git a/dan/ocr/predict/inference.py b/dan/ocr/predict/inference.py index c73c8890318db98098d3ec3894cb35186ce25647..2bea4bee7fb6adf990d87d3b30d66583c009e923 100644 --- a/dan/ocr/predict/inference.py +++ b/dan/ocr/predict/inference.py @@ -59,21 +59,14 @@ class DAN: :param mode: The mode to load the model (train or eval). :param use_language_model: Whether to use an explicit language model to rescore text hypotheses. """ - model_path = list(path.glob("*.pt")) - assert len(model_path) == 1, f"Found {len(model_path)} model(s) `.pt` in {path}" - model_path = model_path.pop() - - params_path = list(path.glob("*parameters.yml")) - assert ( - len(params_path) == 1 - ), f"Found {len(params_path)} parameter(s) `parameters.yml` in {path}" - params_path = params_path.pop() - - charset_path = list(path.glob("*.pkl")) - assert ( - len(charset_path) == 1 - ), f"Found {len(charset_path)} charset(s) `.pkl` in {path}" - charset_path = charset_path.pop() + model_path = path / "model.pt" + assert model_path.is_file(), f"File {model_path} not found" + + params_path = path / "parameters.yml" + assert params_path.is_file(), f"File {params_path} not found" + + charset_path = path / "charset.pkl" + assert charset_path.is_file(), f"File {charset_path} not found" parameters = yaml.safe_load(params_path.read_text())["parameters"] parameters["decoder"]["device"] = self.device diff --git a/docs/usage/predict/index.md b/docs/usage/predict/index.md index 497d9f378840d2f6bf20c9c15e3be1bae76bcd78..9ab4edaad37821e4e42a0edf6ff25cb420f7c99c 100644 --- a/docs/usage/predict/index.md +++ b/docs/usage/predict/index.md @@ -31,9 +31,9 @@ Use the `teklia-dan predict` command to apply a trained DAN model on an image. In the following examples the `models` directory should have: -- a single model file `*.pt`, -- a single charset file `*.pkl`, -- a single parameters file `*parameters.yml`. +- a `model.pt` file, +- a `charset.pkl` file, +- a `parameters.yml` file corresponding to the `inference_parameters.yml` file generated during training. ### Predict with confidence scores diff --git a/tests/data/prediction/popp_line_model.pt b/tests/data/prediction/model.pt similarity index 100% rename from tests/data/prediction/popp_line_model.pt rename to tests/data/prediction/model.pt diff --git a/tests/test_prediction.py b/tests/test_prediction.py index f84b3406b68c02c7ec92e0046f4e9c84577d4908..ffa7bf7005ead1893c5c5d82df43875fded9add9 100644 --- a/tests/test_prediction.py +++ b/tests/test_prediction.py @@ -635,8 +635,8 @@ def test_run_prediction_language_model( model_path.mkdir(exist_ok=True) shutil.copyfile( - PREDICTION_DATA_PATH / "popp_line_model.pt", - model_path / "popp_line_model.pt", + PREDICTION_DATA_PATH / "model.pt", + model_path / "model.pt", ) shutil.copyfile( PREDICTION_DATA_PATH / "charset.pkl",