diff --git a/dan/ocr/predict/prediction.py b/dan/ocr/predict/prediction.py index 064eb749b7675de73653713877f80c9867ffa997..91b4895fde3b8e34a26d9acef6590370b87ff670 100644 --- a/dan/ocr/predict/prediction.py +++ b/dan/ocr/predict/prediction.py @@ -339,7 +339,7 @@ def process_batch( # Return mean confidence score if confidence_score: result["confidences"] = {} - char_confidences = prediction["confidences"][0] + char_confidences = prediction["confidences"][idx] text = result["text"] start_tokens, end_tokens = zip(*list(tokens.values())) end_tokens = list(filter(bool, end_tokens))