Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
  • atr/dan
1 result
Show changes
Commits on Source (2)
Hide whitespace changes
Inline Side-by-side
dan/ocr/predict/inference.py
View file @ 095667f4
...@@ -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}")
......
requirements.txt
View file @ 095667f4
...@@ -10,7 +10,7 @@ numpy==1.24.3 ...@@ -10,7 +10,7 @@ numpy==1.24.3
prettytable==3.8.0 prettytable==3.8.0
PyYAML==6.0 PyYAML==6.0
scipy==1.10.1 scipy==1.10.1
teklia-line-image-extractor==0.2.8rc4 teklia-line-image-extractor==0.2.8rc5
tenacity==8.2.3 tenacity==8.2.3
tensorboard==2.12.2 tensorboard==2.12.2
torch==2.0.0 torch==2.0.0
......
tests/test_prediction.py
View file @ 095667f4
...@@ -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": {},
},
], ],
), ),
), ),
......