Skip to content
Snippets Groups Projects
Verified Commit 9d99164b authored by Mélodie Boillet's avatar Mélodie Boillet
Browse files

Apply d580a769

parent dad76649
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
dan/predict/attention.py
+ 7
− 3
View file @ 9d99164b
......@@ -70,14 +70,18 @@ def split_text_and_confidences(
texts = list(text)
offset = 0
elif level == "word":
texts, probs = compute_prob_by_separator(text, confidences, word_separators)
texts, confidences = compute_prob_by_separator(
text, confidences, word_separators
)
offset = 1
elif level == "line":
texts, probs = compute_prob_by_separator(text, confidences, line_separators)
texts, confidences = compute_prob_by_separator(
text, confidences, line_separators
)
offset = 1
else:
logger.error("Level should be either 'char', 'word', or 'line'")
return texts, [np.around(num, 2) for num in probs], offset
return texts, [np.around(num, 2) for num in confidences], offset
def get_predicted_polygons_with_confidence(
......
dan/predict/prediction.py
+ 2
− 2
View file @ 9d99164b
......@@ -282,9 +282,9 @@ def process_image(
logger.debug("Image pre-processed.")
# Convert to tensor of size (batch_size, channel, height, width) with batch_size=1
input_tensor = im_p.unsqueeze(0)
input_tensor = torch.from_numpy(im_p).permute(2, 0, 1).unsqueeze(0)
input_tensor = input_tensor.to(device)
input_sizes = [im_p.shape[1:]]
input_sizes = [im_p.shape[:2]]
# Parse delimiters to regex
word_separators = parse_delimiters(word_separators)
......
tests/test_prediction.py
+ 227
− 0
View file @ 9d99164b
# -*- coding: utf-8 -*-
import json
import pytest
import torch
from dan.predict.prediction import DAN
from dan.predict.prediction import run as run_prediction
from dan.utils import read_image
......@@ -53,3 +56,227 @@ def test_predict(
prediction = dan_model.predict(input_tensor, input_sizes)
assert prediction == expected_prediction
@pytest.mark.parametrize(
"image_name, confidence_score, temperature, expected_prediction",
(
(
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
None,
1.0,
{"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"},
),
(
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
["word"],
1.0,
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidences": {
"by ner token": [],
"total": 1.0,
"word": [
{"text": "ⓈBellisson", "confidence": 1.0},
{"text": "ⒻGeorges", "confidence": 1.0},
{"text": "Ⓑ91", "confidence": 1.0},
{"text": "ⓁP", "confidence": 1.0},
{"text": "ⒸM", "confidence": 1.0},
{"text": "ⓀCh", "confidence": 1.0},
{"text": "ⓄPlombier", "confidence": 1.0},
{"text": "ⓅPatron?12241", "confidence": 1.0},
],
},
},
),
(
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
["word"],
3.5,
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidences": {
"by ner token": [],
"total": 0.93,
"word": [
{"text": "ⓈBellisson", "confidence": 0.93},
{"text": "ⒻGeorges", "confidence": 0.94},
{"text": "Ⓑ91", "confidence": 0.92},
{"text": "ⓁP", "confidence": 0.94},
{"text": "ⒸM", "confidence": 0.93},
{"text": "ⓀCh", "confidence": 0.96},
{"text": "ⓄPlombier", "confidence": 0.94},
{"text": "ⓅPatron?12241", "confidence": 0.93},
],
},
},
),
(
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
["line"],
1.0,
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidences": {
"by ner token": [],
"total": 1.0,
"line": [
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidence": 1.0,
}
],
},
},
),
(
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
["line"],
3.5,
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidences": {
"by ner token": [],
"total": 0.93,
"line": [
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"confidence": 0.93,
}
],
},
},
),
(
"0dfe8bcd-ed0b-453e-bf19-cc697012296e",
None,
1.0,
{"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376"},
),
(
"0dfe8bcd-ed0b-453e-bf19-cc697012296e",
["char", "word", "line"],
1.0,
{
"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
"confidences": {
"by ner token": [],
"total": 1.0,
"char": [
{"text": "Ⓢ", "confidence": 1.0},
{"text": "T", "confidence": 1.0},
{"text": "e", "confidence": 1.0},
{"text": "m", "confidence": 1.0},
{"text": "p", "confidence": 1.0},
{"text": "l", "confidence": 1.0},
{"text": "i", "confidence": 1.0},
{"text": "é", "confidence": 0.85},
{"text": " ", "confidence": 1.0},
{"text": "Ⓕ", "confidence": 1.0},
{"text": "M", "confidence": 1.0},
{"text": "a", "confidence": 1.0},
{"text": "r", "confidence": 1.0},
{"text": "c", "confidence": 1.0},
{"text": "e", "confidence": 1.0},
{"text": "l", "confidence": 1.0},
{"text": "l", "confidence": 1.0},
{"text": "e", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "Ⓑ", "confidence": 1.0},
{"text": "9", "confidence": 1.0},
{"text": "3", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "Ⓛ", "confidence": 1.0},
{"text": "S", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "Ⓚ", "confidence": 1.0},
{"text": "c", "confidence": 1.0},
{"text": "h", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "Ⓞ", "confidence": 1.0},
{"text": "E", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "d", "confidence": 1.0},
{"text": "a", "confidence": 1.0},
{"text": "c", "confidence": 1.0},
{"text": "t", "confidence": 1.0},
{"text": "y", "confidence": 1.0},
{"text": "l", "confidence": 1.0},
{"text": "o", "confidence": 1.0},
{"text": " ", "confidence": 1.0},
{"text": "Ⓟ", "confidence": 1.0},
{"text": "1", "confidence": 1.0},
{"text": "8", "confidence": 1.0},
{"text": "3", "confidence": 1.0},
{"text": "7", "confidence": 1.0},
{"text": "6", "confidence": 1.0},
],
"word": [
{"text": "ⓈTemplié", "confidence": 0.98},
{"text": "ⒻMarcelle", "confidence": 1.0},
{"text": "Ⓑ93", "confidence": 1.0},
{"text": "ⓁS", "confidence": 1.0},
{"text": "Ⓚch", "confidence": 1.0},
{"text": "ⓄE", "confidence": 1.0},
{"text": "dactylo", "confidence": 1.0},
{"text": "Ⓟ18376", "confidence": 1.0},
],
"line": [
{
"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
"confidence": 1.0,
}
],
},
},
),
(
"2c242f5c-e979-43c4-b6f2-a6d4815b651d",
False,
1.0,
{"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31"},
),
(
"ffdec445-7f14-4f5f-be44-68d0844d0df1",
False,
1.0,
{"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"},
),
),
)
def test_run_prediction(
image_name,
confidence_score,
temperature,
expected_prediction,
prediction_data_path,
tmp_path,
):
run_prediction(
image=(prediction_data_path / "images" / image_name).with_suffix(".png"),
image_dir=None,
model=prediction_data_path / "popp_line_model.pt",
parameters=prediction_data_path / "parameters.yml",
charset=prediction_data_path / "charset.pkl",
output=tmp_path,
scale=1,
confidence_score=True if confidence_score else False,
confidence_score_levels=confidence_score if confidence_score else [],
attention_map=False,
attention_map_level=None,
attention_map_scale=0.5,
word_separators=[" ", "\n"],
line_separators=["\n"],
temperature=temperature,
image_max_width=None,
predict_objects=False,
threshold_method="otsu",
threshold_value=0,
image_extension=None,
gpu_device=None,
)
with (tmp_path / image_name).with_suffix(".json").open("r") as json_file:
prediction = json.load(json_file)
assert prediction == expected_prediction
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment