Load the model via a path to a folder

Closes #210 (closed)

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