From 6c95feece66642e75e1cd1befa6cbab4a57a19c4 Mon Sep 17 00:00:00 2001 From: Kohulan Rajan Date: Wed, 30 Aug 2023 09:14:25 +0200 Subject: [PATCH] fix: confidence extraction --- DECIMER/Repack_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DECIMER/Repack_model.py b/DECIMER/Repack_model.py index aa61f00..1d94ead 100644 --- a/DECIMER/Repack_model.py +++ b/DECIMER/Repack_model.py @@ -190,7 +190,7 @@ def __call__(self, Decoded_image): predictions = prediction_batch[:, -1:, :] # (batch_size, 1, vocab_size) predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32) - confidence = predictions[-1][-1][int(predicted_id)] + confidence = predictions[0, 0, int(predicted_id[0, 0])] output_array = output_array.write(t + 1, predicted_id[0]) confidence_array = confidence_array.write(t + 1, confidence) if predicted_id == end_token: