diff --git a/dan/datasets/extract/utils.py b/dan/datasets/extract/utils.py
index f9c2224528938f4953b84d7443c4db7ae0ffce50..2b93334c8320594b8251497f7e36fa4b8a6255f6 100644
--- a/dan/datasets/extract/utils.py
+++ b/dan/datasets/extract/utils.py
@@ -87,16 +87,18 @@ def insert_token(text: str, entity_type: EntityType, offset: int, length: int) -
def normalize_linebreaks(text: str) -> str:
"""
- Remove begin/ending linebreaks
- Replace \r with regular linebreak and consecutive linebreaks
+ Remove begin/ending linebreaks.
+ Replace \r with regular linebreak and consecutive linebreaks.
+ :param text: Text to normalize.
"""
return TRIM_RETURN_REGEX.sub("\n", text.strip())
def normalize_spaces(text: str) -> str:
"""
- Remove begin/ending spaces
- Replace \t with regular space and consecutive spaces
+ Remove begin/ending spaces.
+ Replace \t with regular space and consecutive spaces.
+ :param text: Text to normalize.
"""
return TRIM_SPACE_REGEX.sub(" ", text.strip())
diff --git a/dan/ocr/decoder.py b/dan/ocr/decoder.py
index b4da94d0464c37d6543b8b7310faa5fbed06776b..b62c0f6c77b48518dd0939e984c186b1377e887d 100644
--- a/dan/ocr/decoder.py
+++ b/dan/ocr/decoder.py
@@ -470,17 +470,13 @@ class GlobalHTADecoder(Module):
class CTCLanguageDecoder:
"""
Initialize a CTC decoder with n-gram language modeling.
- Args:
- language_model_path (str): path to a KenLM or ARPA language model
- lexicon_path (str): path to a lexicon file containing the possible words and corresponding spellings.
- Each line consists of a word and its space separated spelling. If `None`, uses lexicon-free
- decoding.
- tokens_path (str): path to a file containing valid tokens. If using a file, the expected
- format is for tokens mapping to the same index to be on the same line
- language_model_weight (float): weight of the language model.
- blank_token (str): token representing the blank/ctc symbol
- unk_token (str): token representing unknown characters
- sil_token (str): token representing the space character
+ :param language_model_path: Path to a KenLM or ARPA language model.
+ :param lexicon_path: Path to a lexicon file containing the possible words and corresponding spellings.
+ Each line consists of a word and its space separated spelling. If `None`, uses lexicon-free decoding.
+ :param tokens_path: Path to a file containing valid tokens. If using a file, the expected
+ format is for tokens mapping to the same index to be on the same line.
+ :param language_model_weight: Weight of the language model.
+ :param temperature: Temperature for model calibreation.
"""
def __init__(
@@ -499,6 +495,9 @@ class CTCLanguageDecoder:
}
self.index_to_token = {i: token for token, i in self.tokens_to_index.items()}
self.blank_token_id = self.tokens_to_index[self.mapping.ctc.encoded]
+
+ # Torchaudio's decoder
+ # https://pytorch.org/audio/master/generated/torchaudio.models.decoder.ctc_decoder.html
self.decoder = ctc_decoder(
lm=language_model_path,
lexicon=lexicon_path,
@@ -516,7 +515,7 @@ class CTCLanguageDecoder:
self, batch_features: torch.FloatTensor, batch_frames: torch.LongTensor
) -> tuple[torch.FloatTensor, torch.LongTensor]:
"""
- Add CTC frames between each characters to avoid duplicate removal
+ Add CTC frames between each characters to avoid duplicate removal.
"""
high_prob = batch_features.max()
low_prob = batch_features.min()
@@ -563,6 +562,9 @@ class CTCLanguageDecoder:
) -> Dict[str, List[Union[str, float]]]:
"""
Post-process hypotheses to output JSON. Exports only the best hypothesis for each image.
+ :param hypotheses: List of hypotheses returned by the decoder.
+ :param batch_sizes: Prediction length of size batch_size.
+ :return: A dictionary containing the hypotheses and their confidences.
"""
out = {}
# Replace <space> by an actual space and format string
@@ -595,11 +597,9 @@ class CTCLanguageDecoder:
) -> Dict[str, List[Union[str, float]]]:
"""
Decode a feature vector using n-gram language modelling.
- Args:
- features: Feature vector of size (batch_size, n_tokens, n_frames).
- batch_sizes: Prediction length of size (batch_size)
- Returns:
- a dictionary containing the hypotheses and their confidences
+ :param batch_features: Feature vector of size (batch_size, n_tokens, n_frames).
+ :param batch_frames: Prediction length of size batch_size.
+ :return: A dictionary containing the hypotheses and their confidences.
"""
# Reshape from (batch_size, n_tokens, n_frames) to (batch_size, n_frames, n_tokens)
batch_features = batch_features.permute((0, 2, 1))
diff --git a/dan/ocr/predict/prediction.py b/dan/ocr/predict/prediction.py
index 4ce2898f3c7ce8b6b6b8b481adeececa59466831..0719f4926ef8d7497bf491833799cb3f1b4a75d3 100644
--- a/dan/ocr/predict/prediction.py
+++ b/dan/ocr/predict/prediction.py
@@ -479,7 +479,6 @@ def run(
dan_model.load(
model, parameters, charset, mode="eval", use_language_model=use_language_model
)
- batch_size = 1 if use_language_model else batch_size
# Do not use LM with invalid LM weight
use_language_model = dan_model.lm_decoder is not None
diff --git a/dan/utils.py b/dan/utils.py
index e86e08f59c1c9ccb7b44265616287fe7f2baf472..293974b98ed38d46e98ff43fa821b26b3feaaf97 100644
--- a/dan/utils.py
+++ b/dan/utils.py
@@ -139,7 +139,9 @@ def parse_tokens(filename: str) -> Dict[str, EntityType]:
def read_yaml(yaml_path: str) -> Dict:
"""
- Read YAML tokens file
+ Read YAML tokens file.
+ :param yaml_path: Path of the YAML file to read.
+ :return: The content of the read file.
"""
filename = Path(yaml_path)
assert filename.exists(), f"{yaml_path} does not resolve."
@@ -152,6 +154,8 @@ def read_yaml(yaml_path: str) -> Dict:
def read_json(json_path: str) -> Dict:
"""
Read labels JSON file
+ :param json_path: Path of the JSON file to read.
+ :return: The content of the read file.
"""
filename = Path(json_path)
assert filename.exists(), f"{json_path} does not resolve."
@@ -163,7 +167,9 @@ def read_json(json_path: str) -> Dict:
def read_txt(txt_path: str) -> str:
"""
- Read TXT file
+ Read TXT file.
+ :param txt_path: Path of the text file to read.
+ :return: The content of the read file.
"""
filename = Path(txt_path)
assert filename.exists(), f"{txt_path} does not resolve."
diff --git a/docs/get_started/training.md b/docs/get_started/training.md
index 03773c50db6d9ccff6a0abdfa3167d47f201f373..1cb5a0cffb55d6aaf9c9e2fee9fc65e433add1af 100644
--- a/docs/get_started/training.md
+++ b/docs/get_started/training.md
@@ -21,9 +21,9 @@ output/
│ ├── val
│ └── test
├── language_model
-│ ├── language_corpus.txt
-│ ├── language_lexicon.txt
-│ └── language_tokens.txt
+│ ├── corpus.txt
+│ ├── lexicon.txt
+│ └── tokens.txt
```
## 2. Train
diff --git a/docs/usage/predict.md b/docs/usage/predict.md
deleted file mode 100644
index 4a13bd227a38ddc347670c522524908782b2a519..0000000000000000000000000000000000000000
--- a/docs/usage/predict.md
+++ /dev/null
@@ -1,214 +0,0 @@
-# Predict
-
-Use the `teklia-dan predict` command to apply a trained DAN model on an image.
-
-## Description of parameters
-
-| Parameter | Description | Type | Default |
-| --------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ------------- |
-| `--image` | Path to the image to predict. Must not be provided with `--image-dir`. | `Path` | |
-| `--image-dir` | Path to the folder where the images to predict are stored. Must not be provided with `--image`. | `Path` | |
-| `--image-extension` | The extension of the images in the folder. Ignored if `--image-dir` is not provided. | `str` | .jpg |
-| `--model` | Path to the model to use for prediction | `Path` | |
-| `--parameters` | Path to the YAML parameters file. | `Path` | |
-| `--charset` | Path to the charset file. | `Path` | |
-| `--output` | Path to the output folder. Results will be saved in this directory. | `Path` | |
-| `--confidence-score` | Whether to return confidence scores. | `bool` | `False` |
-| `--confidence-score-levels` | Level to return confidence scores. Should be any combination of `["line", "word", "char"]`. | `str` | |
-| `--attention-map` | Whether to plot attention maps. | `bool` | `False` |
-| `--attention-map-scale` | Image scaling factor before creating the GIF. | `float` | `0.5` |
-| `--attention-map-level` | Level to plot the attention maps. Should be in `["line", "word", "char"]`. | `str` | `"line"` |
-| `--predict-objects` | Whether to return polygons coordinates. | `bool` | `False` |
-| `--word-separators` | List of word separators. | `list` | `[" ", "\n"]` |
-| `--line-separators` | List of line separators. | `list` | `["\n"]` |
-| `--threshold-method` | Method to use for attention mask thresholding. Should be in `["otsu", "simple"]`. | `str` | `"otsu"` |
-| `--threshold-value ` | Threshold to use for the "simple" thresholding method. | `int` | `0` |
-| `--batch-size ` | Size of the batches for prediction. | `int` | `1` |
-| `--start-token ` | Use a specific starting token at the beginning of the prediction. Useful when making predictions on different single pages. | `str` | `None` |
-| `--use-language-model` | Whether to use an external n-gram language model to rescore hypotheses. See [the dedicated example](#predict-with-an-external-n-gram-language-model) for details. | `bool` | `False` |
-
-## Examples
-
-### Predict with confidence scores
-
-To run a prediction with confidence scores, run this command:
-
-```shell
-teklia-dan predict \
- --image dan_humu_page/example.jpg \
- --model dan_humu_page/model.pt \
- --parameters dan_humu_page/parameters.yml \
- --charset dan_humu_page/charset.pkl \
- --output dan_humu_page/predict/ \
- --confidence-score
-```
-
-It will create the following JSON file named `dan_humu_page/predict/example.json`
-
-```json
-{
- "text": "Hansteensgt. 2 IV 28/4 - 19\nKj\u00e6re Gerhard.\nTak for Brevet om Boken og Haven\nog Crokus og Blaaveis og tak fordi\nDu vilde be mig derut sammen\nmed Kris og Ragna. Men vet Du\nda ikke, at Kris reiste med sin S\u00f8-\nster Fru Cr\u00f8ger til Lillehammer\nnogle Dage efter Begravelsen? Hen\ndes Address er Amtsingeni\u00f8r\nCr\u00f8ger. Hun skriver at de blir\nder til lidt ut i Mai. Nu er hun\nnoksaa medtat skj\u00f8nner jeg af Sorg\nog af L\u00e6ngsel, skriver saameget r\u00f8-\nrende om Oluf. Ragna har det\nherligt, skriver hun. Hun er bare\ngla, og det vet jeg, at \"Oluf er gla over,\nder hvor han nu er. Jeg har saa in-\nderlig ondt af hende, og om Du skrev\net Par Ord tror jeg det vilde gj\u00f8re\nhende godt. - Jeg gl\u00e6der mig over,\nat Du har skrevet en Bok, og\njeg er vis paa, at den er god.",
- "confidence": 0.99
-}
-```
-
-### Predict with confidence scores and line-level attention maps
-
-To run a prediction with confidence scores and plot line-level attention maps, run this command:
-
-```shell
-teklia-dan predict \
- --image dan_humu_page/example.jpg \
- --model dan_humu_page/model.pt \
- --parameters dan_humu_page/parameters.yml \
- --charset dan_humu_page/charset.pkl \
- --output dan_humu_page/predict/ \
- --confidence-score \
- --attention-map \
-```
-
-It will create the following JSON file named `dan_humu_page/predict/example.json` and a GIF showing a word-level attention map `dan_humu_page/predict/example_line.gif`
-
-```json
-{
- "text": "Hansteensgt. 2 IV 28/4 - 19\nKj\u00e6re Gerhard.\nTak for Brevet om Boken og Haven\nog Crokus og Blaaveis og tak fordi\nDu vilde be mig derut sammen\nmed Kris og Ragna. Men vet Du\nda ikke, at Kris reiste med sin S\u00f8-\nster Fru Cr\u00f8ger til Lillehammer\nnogle Dage efter Begravelsen? Hen\ndes Address er Amtsingeni\u00f8r\nCr\u00f8ger. Hun skriver at de blir\nder til lidt ut i Mai. Nu er hun\nnoksaa medtat skj\u00f8nner jeg af Sorg\nog af L\u00e6ngsel, skriver saameget r\u00f8-\nrende om Oluf. Ragna har det\nherligt, skriver hun. Hun er bare\ngla, og det vet jeg, at \"Oluf er gla over,\nder hvor han nu er. Jeg har saa in-\nderlig ondt af hende, og om Du skrev\net Par Ord tror jeg det vilde gj\u00f8re\nhende godt. - Jeg gl\u00e6der mig over,\nat Du har skrevet en Bok, og\njeg er vis paa, at den er god.",
- "confidence": 0.99,
- "attention_gif": "dan_humu_page/predict/example_line.gif"
-}
-```
-
-<img src="../../assets/example_line.gif" />
-
-### Predict with confidence scores and word-level attention maps
-
-To run a prediction with confidence scores and plot word-level attention maps, run this command:
-
-```shell
-teklia-dan predict \
- --image dan_humu_page/example.jpg \
- --model dan_humu_page/model.pt \
- --parameters dan_humu_page/parameters.yml \
- --charset dan_humu_page/charset.pkl \
- --output dan_humu_page/predict/ \
- --confidence-score \
- --attention-map \
- --attention-map-level word \
- --attention-map-scale 0.5
-```
-
-It will create the following JSON file named `dan_humu_page/predict/example.json` and a GIF showing a word-level attention map `dan_humu_page/predict/example_word.gif`.
-
-```json
-{
- "text": "Hansteensgt. 2 IV 28/4 - 19\nKj\u00e6re Gerhard.\nTak for Brevet om Boken og Haven\nog Crokus og Blaaveis og tak fordi\nDu vilde be mig derut sammen\nmed Kris og Ragna. Men vet Du\nda ikke, at Kris reiste med sin S\u00f8-\nster Fru Cr\u00f8ger til Lillehammer\nnogle Dage efter Begravelsen? Hen\ndes Address er Amtsingeni\u00f8r\nCr\u00f8ger. Hun skriver at de blir\nder til lidt ut i Mai. Nu er hun\nnoksaa medtat skj\u00f8nner jeg af Sorg\nog af L\u00e6ngsel, skriver saameget r\u00f8-\nrende om Oluf. Ragna har det\nherligt, skriver hun. Hun er bare\ngla, og det vet jeg, at \"Oluf er gla over,\nder hvor han nu er. Jeg har saa in-\nderlig ondt af hende, og om Du skrev\net Par Ord tror jeg det vilde gj\u00f8re\nhende godt. - Jeg gl\u00e6der mig over,\nat Du har skrevet en Bok, og\njeg er vis paa, at den er god.",
- "confidence": 0.99,
- "attention_gif": "dan_humu_page/predict/example_word.gif"
-}
-```
-
-<img src="../../assets/example_word.gif" >
-
-### Predict with line-level attention maps and extract polygons
-
-To run a prediction, plot line-level attention maps, and extract polygons, run this command:
-
-```shell
-teklia-dan predict \
- --image dan_humu_page/example.jpg \
- --model dan_humu_page/model.pt \
- --parameters dan_humu_page/parameters.yml \
- --charset dan_humu_page/charset.pkl \
- --output dan_humu_page/predict/ \
- --attention-map \
- --predict-objects \
- --threshold-method otsu
-```
-
-It will create the following JSON file named `dan_humu_page/predict/example.json` and a GIF showing a line-level attention map with extracted polygons `dan_humu_page/predict/example_line.gif`
-
-```json
-{
- "text": "Oslo\n39 \nOresden den 24te Rasser!\nH\u00f8jst\u00e6redesherr Hartvig - assert!\nUllereder fra den f\u00f8rste tide da\njeg havder den tilfredsstillelser at vide den ar-\ndistiske ledelser af Kristiania theater i Deres\nhronder, har jeg g\u00e5t hernede med et stille\nh\u00e5b om fra Dem at modtage et forelag, sig -\nsende tils at lade \"K\u00e6rlighedens \u00abKomedie\u00bb\nopf\u00f8re fore det norske purblikum.\nEt s\u00e5dant forslag er imidlertid, imod\nforventning; ikke fremkommet, og jeg n\u00f8des der-\nfor tils self at grivbe initiativet, hvilket hervede\nsker, idet jeg\nbeder\nbet\nragte stigkket some ved denne\nskrivelse officielde indleveret til theatret. No-\nget exemplar af bogen vedlagger jeg ikke da\ndenne (i 2den udgave) med Lethed kan er -\nholdet deroppe.\nDe bet\u00e6nkeligheder, jeg i sin tid n\u00e6-\nrede mod stykkets opf\u00f8relse, er for l\u00e6nge si -\ndem forsvundne. Af mange begn er jeg kom-\nmen til den overbevisning at almenlreden\naru har f\u00e5tt sine \u00f8gne opladte for den sand -\nMed at dette arbejde i sin indersten id\u00e9 hviler\np\u00e5 et ubedinget meralsk grundlag, og brad\nstykkets hele kunstneriske struktuve ang\u00e5r,",
- "objects": [
- {
- "confidence": 0.68,
- "polygon": [
- [
- 264,
- 118
- ],
- [
- 410,
- 118
- ],
- [
- 410,
- 185
- ],
- [
- 264,
- 185
- ]
- ],
- "text": "Oslo",
- "text_confidence": 0.8
- }
- ],
- "attention_gif": "dan_humu_page/predict/example_line.gif"
-}
-```
-
-<img src="../../assets/example_line_polygon.gif" >
-
-### Predict with an external n-gram language model
-
-#### Build the language model
-
-A dataset extracted with the `teklia-dan dataset extract` command should contain the files required to build a language model (in the `language_model` folder). To refine DAN's predictions with a language model, follow these steps:
-
-1. Install and build [kenlm](https://github.com/kpu/kenlm)
-1. Build a 6-gram language model using the following command
-
-```sh
-bin/lmplz --order 6 \
- --text my_dataset/language_model/corpus.txt \
- --arpa my_dataset/language_model/model.arpa
-```
-
-1. Update `inference_parameters.yml`. The `weight` parameter defines how much weight to give to the language model. It should be set carefully (usually between 0.5 and 2.0) as it will affect the quality of the predictions.
-
-```yaml
-parameters:
- ...
- language_model:
- model: my_dataset/language_model/model.arpa
- lexicon: my_dataset/language_model/lexicon.txt
- tokens: my_dataset/language_model/tokens.txt
- weight: 0.5
-```
-
-#### Predict
-
-To run a prediction with the n-gram language model, run this command:
-
-```shell
-teklia-dan predict \
- --image dan_humu_page/example.jpg \
- --model dan_humu_page/model.pt \
- --parameters dan_humu_page/parameters.yml \
- --charset dan_humu_page/charset.pkl \
- --use-language-model \
- --output dan_humu_page/predict/
-```
-
-It will create the following JSON file named `dan_humu_page/predict/example.json`
-
-```json
-{
- "text": "etc., some jeg netop idag\nholder Vask paa.\nLeien af Skj\u00f8rterne\nbestad i at jeg kj\u00f8bte\net Forkl\u00e6de til hver\naf de to Piger, some\nhavde laant os dem.\nResten var Vask af Hardan-\ngerskj\u00f8rter og et Forkl\u00e6de,\nsamt Fragt paa det Gods\n(N\u00f8i) some man sendte\nmig ubet\u00e6lt.\nIdag fik jeg hyggeligt\nFrimarkebrev fra Fosvold\nMed Hilsen\nDeres\nHulda Garborg",
- "language_model": {
- "text": "eet., some jeg netop idag\nholder Vask paa.\nLeien af Skj\u00f9rterne\nbestad i at jeg kj\u00f9bte\net Forkl\u00e7de til hver\naf de to Piger, some\nhavde laant os dem.\nResten var Vask af Hardan-\ngerskj\u00f9rter og et Forkl\u00e7de,\nsamt Fragt paa det Gods\n(N\u00f9i) some man sendte\nmig ubetalt.\nIdag fik jeg hyggeligt\nFrimarkebrev fra Fosvold\nMed Hilsen\nDeres\nHulda Garborg",
- "confidence": 0.87
- }
-}
-```
diff --git a/tests/test_extract.py b/tests/test_extract.py
index 49841521d928d18a6c78c668a5d7cab97574c3f8..c9b4d3dea6bedb5756f7d7ede5325151c6b8809b 100644
--- a/tests/test_extract.py
+++ b/tests/test_extract.py
@@ -318,11 +318,11 @@ def test_process_element_unknown_token_in_text_error(mock_database, tmp_path):
arkindex_extractor.process_element(element, "val")
-@pytest.mark.parametrize("load_entities", (True,)) # False))
-@pytest.mark.parametrize("keep_spaces", (True,)) # False))
+@pytest.mark.parametrize("load_entities", (True, False))
+@pytest.mark.parametrize("keep_spaces", (True, False))
# Transcription and entities have the same worker version
@pytest.mark.parametrize(
- "transcription_entities_worker_version", (False,) # "worker_version_id",)#, False)
+ "transcription_entities_worker_version", ("worker_version_id", False)
)
@patch("dan.datasets.extract.extract.download_image")
def test_extract(
diff --git a/tests/test_prediction.py b/tests/test_prediction.py
index b3a1bf578468a23f6977993eb897f6a25c217423..cc61d2ecf43a271740516f87a0c5bb3e719f74a1 100644
--- a/tests/test_prediction.py
+++ b/tests/test_prediction.py
@@ -589,12 +589,10 @@ def test_run_prediction_batch(
),
),
)
-@pytest.mark.parametrize("batch_size", [1, 2])
def test_run_prediction_language_model(
image_names,
language_model_weight,
expected_predictions,
- batch_size,
tmp_path,
):
# Make tmpdir and copy needed images inside