From f000187f1d718fb792bb2ccb9e1cc20d2e394bc8 Mon Sep 17 00:00:00 2001 From: Moses Paul R Date: Fri, 13 Dec 2024 17:37:47 +0000 Subject: [PATCH] fix missing confidence in surya layout predictions --- surya/layout.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/surya/layout.py b/surya/layout.py index c30d6f6..476e54e 100644 --- a/surya/layout.py +++ b/surya/layout.py @@ -219,11 +219,14 @@ def batch_layout_detection(images: List, model, processor, batch_size=None, top_ top_k_indices = [p["top_k_indices"] - model.decoder.config.special_token_count for p in preds] for z, (poly, label) in enumerate(zip(polygons, labels)): + l = ID_TO_LABEL[int(label)] + top_k_dict = {ID_TO_LABEL.get(int(l)): prob.item() for (l, prob) in zip(top_k_indices[z], top_k_probs[z]) if l > 0} lb = LayoutBox( polygon=poly, - label=ID_TO_LABEL[int(label)], + label=l, position=z, - top_k={ID_TO_LABEL.get(int(l)): prob.item() for (l, prob) in zip(top_k_indices[z], top_k_probs[z]) if l > 0} + top_k=top_k_dict, + confidence=top_k_dict[l] ) boxes.append(lb) boxes = clean_boxes(boxes)