Skip to content
Snippets Groups Projects
Commit deb015a0 authored by Yoann Schneider's avatar Yoann Schneider :tennis:
Browse files

Support the case where DAN does not output anything

parent 307df1b7
No related branches found
No related tags found
1 merge request!298Support the case where DAN does not output anything
Hide whitespace changes
Inline Side-by-side
dan/ocr/predict/inference.py
+ 54
− 54
View file @ deb015a0
...@@ -356,62 +356,62 @@ def process_batch( ...@@ -356,62 +356,62 @@ def process_batch(
logger.info("Prediction parsing...") logger.info("Prediction parsing...")
for idx, image_path in enumerate(image_batch): for idx, image_path in enumerate(image_batch):
predicted_text = prediction["text"][idx] predicted_text = prediction["text"][idx]
result = {"text": predicted_text} result = {"text": predicted_text, "confidences": {}, "language_model": {}}
# Return LM results if predicted_text:
if use_language_model: # Return LM results
result["language_model"] = { if use_language_model:
"text": prediction["language_model"]["text"][idx], result["language_model"] = {
"confidence": prediction["language_model"]["confidence"][idx], "text": prediction["language_model"]["text"][idx],
} "confidence": prediction["language_model"]["confidence"][idx],
}
# Return extracted objects (coordinates, text, confidence)
if predict_objects: # Return extracted objects (coordinates, text, confidence)
result["objects"] = prediction["objects"][idx] if predict_objects:
result["objects"] = prediction["objects"][idx]
# Return mean confidence score
if confidence_score: # Return mean confidence score
result["confidences"] = {} if confidence_score:
char_confidences = prediction["confidences"][idx] char_confidences = prediction["confidences"][idx]
result["confidences"]["total"] = np.around(np.mean(char_confidences), 2) result["confidences"]["total"] = np.around(np.mean(char_confidences), 2)
for level in confidence_score_levels: for level in confidence_score_levels:
result["confidences"][level.value] = [] result["confidences"][level.value] = []
texts, confidences, _ = split_text_and_confidences( texts, confidences, _ = split_text_and_confidences(
predicted_text, predicted_text,
char_confidences, char_confidences,
level, level,
word_separators, word_separators,
line_separators, line_separators,
tokens, tokens,
)
for text, conf in zip(texts, confidences):
result["confidences"][level.value].append(
{"text": text, "confidence": conf}
) )
# Save gif with attention map for text, conf in zip(texts, confidences):
if attention_map: result["confidences"][level.value].append(
attentions = prediction["attentions"][idx] {"text": text, "confidence": conf}
gif_filename = f"{output}/{image_path.stem}_{attention_map_level}.gif" )
logger.info(f"Creating attention GIF in {gif_filename}")
plot_attention( # Save gif with attention map
image=visu_tensor[idx], if attention_map:
text=predicted_text, attentions = prediction["attentions"][idx]
weights=attentions, gif_filename = f"{output}/{image_path.stem}_{attention_map_level}.gif"
level=attention_map_level, logger.info(f"Creating attention GIF in {gif_filename}")
scale=attention_map_scale, plot_attention(
word_separators=word_separators, image=visu_tensor[idx],
line_separators=line_separators, text=predicted_text,
tokens=tokens, weights=attentions,
display_polygons=predict_objects, level=attention_map_level,
threshold_method=threshold_method, scale=attention_map_scale,
threshold_value=threshold_value, word_separators=word_separators,
max_object_height=max_object_height, line_separators=line_separators,
outname=gif_filename, tokens=tokens,
) display_polygons=predict_objects,
result["attention_gif"] = gif_filename threshold_method=threshold_method,
threshold_value=threshold_value,
max_object_height=max_object_height,
outname=gif_filename,
)
result["attention_gif"] = gif_filename
json_filename = Path(output, image_path.stem).with_suffix(".json") json_filename = Path(output, image_path.stem).with_suffix(".json")
logger.info(f"Saving JSON prediction in {json_filename}") logger.info(f"Saving JSON prediction in {json_filename}")
......
tests/test_prediction.py
+ 46
− 7
View file @ deb015a0
...@@ -67,7 +67,11 @@ def test_predict(image_name, expected_prediction): ...@@ -67,7 +67,11 @@ def test_predict(image_name, expected_prediction):
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84", "0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
None, None,
1.0, 1.0,
{"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"}, {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": {},
},
), ),
( (
"0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84", "0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84",
...@@ -75,6 +79,7 @@ def test_predict(image_name, expected_prediction): ...@@ -75,6 +79,7 @@ def test_predict(image_name, expected_prediction):
1.0, 1.0,
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"word": [ "word": [
...@@ -96,6 +101,7 @@ def test_predict(image_name, expected_prediction): ...@@ -96,6 +101,7 @@ def test_predict(image_name, expected_prediction):
3.5, 3.5,
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 0.93, "total": 0.93,
"ner": [ "ner": [
...@@ -127,6 +133,7 @@ def test_predict(image_name, expected_prediction): ...@@ -127,6 +133,7 @@ def test_predict(image_name, expected_prediction):
1.0, 1.0,
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"line": [ "line": [
...@@ -144,6 +151,7 @@ def test_predict(image_name, expected_prediction): ...@@ -144,6 +151,7 @@ def test_predict(image_name, expected_prediction):
3.5, 3.5,
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 0.93, "total": 0.93,
"ner": [ "ner": [
...@@ -169,7 +177,11 @@ def test_predict(image_name, expected_prediction): ...@@ -169,7 +177,11 @@ def test_predict(image_name, expected_prediction):
"0dfe8bcd-ed0b-453e-bf19-cc697012296e", "0dfe8bcd-ed0b-453e-bf19-cc697012296e",
None, None,
1.0, 1.0,
{"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376"}, {
"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
"language_model": {},
"confidences": {},
},
), ),
( (
"0dfe8bcd-ed0b-453e-bf19-cc697012296e", "0dfe8bcd-ed0b-453e-bf19-cc697012296e",
...@@ -177,6 +189,7 @@ def test_predict(image_name, expected_prediction): ...@@ -177,6 +189,7 @@ def test_predict(image_name, expected_prediction):
1.0, 1.0,
{ {
"text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376", "text": "ⓈTemplié ⒻMarcelle Ⓑ93 ⓁS Ⓚch ⓄE dactylo Ⓟ18376",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"ner": [ "ner": [
...@@ -260,13 +273,21 @@ def test_predict(image_name, expected_prediction): ...@@ -260,13 +273,21 @@ def test_predict(image_name, expected_prediction):
"2c242f5c-e979-43c4-b6f2-a6d4815b651d", "2c242f5c-e979-43c4-b6f2-a6d4815b651d",
False, False,
1.0, 1.0,
{"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31"}, {
"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31",
"language_model": {},
"confidences": {},
},
), ),
( (
"ffdec445-7f14-4f5f-be44-68d0844d0df1", "ffdec445-7f14-4f5f-be44-68d0844d0df1",
False, False,
1.0, 1.0,
{"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"}, {
"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
"language_model": {},
"confidences": {},
},
), ),
), ),
) )
...@@ -315,7 +336,13 @@ def test_run_prediction( ...@@ -315,7 +336,13 @@ def test_run_prediction(
["0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84"], ["0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84"],
None, None,
1.0, 1.0,
[{"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241"}], [
{
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": {},
}
],
), ),
( (
["0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84"], ["0a56e8b3-95cd-4fa5-a17b-5b0ff9e6ea84"],
...@@ -324,6 +351,7 @@ def test_run_prediction( ...@@ -324,6 +351,7 @@ def test_run_prediction(
[ [
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"word": [ "word": [
...@@ -350,6 +378,7 @@ def test_run_prediction( ...@@ -350,6 +378,7 @@ def test_run_prediction(
[ [
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"ner": [ "ner": [
...@@ -376,6 +405,7 @@ def test_run_prediction( ...@@ -376,6 +405,7 @@ def test_run_prediction(
}, },
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"ner": [ "ner": [
...@@ -409,6 +439,7 @@ def test_run_prediction( ...@@ -409,6 +439,7 @@ def test_run_prediction(
[ [
{ {
"text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241", "text": "ⓈBellisson ⒻGeorges Ⓑ91 ⓁP ⒸM ⓀCh ⓄPlombier ⓅPatron?12241",
"language_model": {},
"confidences": { "confidences": {
"total": 1.0, "total": 1.0,
"word": [ "word": [
...@@ -433,8 +464,16 @@ def test_run_prediction( ...@@ -433,8 +464,16 @@ def test_run_prediction(
False, False,
1.0, 1.0,
[ [
{"text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31"}, {
{"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère"}, "text": "Ⓢd ⒻCharles Ⓑ11 ⓁP ⒸC ⓀF Ⓞd Ⓟ14 31",
"language_model": {},
"confidences": {},
},
{
"text": "ⓈNaudin ⒻMarie Ⓑ53 ⓁS Ⓒv ⓀBelle mère",
"language_model": {},
"confidences": {},
},
], ],
), ),
), ),
......
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