fix: KeyError in tableformer prediction (#854)

* fix for KeyError in tableformer prediction

Signed-off-by: Maksym Lysak <mly@zurich.ibm.com>

* chore: rewrite cumbersome dictionary checking

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>

---------

Signed-off-by: Maksym Lysak <mly@zurich.ibm.com>
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
Co-authored-by: Maksym Lysak <mly@zurich.ibm.com>
Co-authored-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Maxim Lysak 2025-01-31 17:00:14 +01:00 committed by GitHub
parent 70d68b6164
commit b1cf796730
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -209,12 +209,16 @@ class TableStructureModel(BasePageModel):
tc.bbox = tc.bbox.scaled(1 / self.scale)
table_cells.append(tc)
assert "predict_details" in table_out
# Retrieving cols/rows, after post processing:
num_rows = table_out["predict_details"]["num_rows"]
num_cols = table_out["predict_details"]["num_cols"]
otsl_seq = table_out["predict_details"]["prediction"][
"rs_seq"
]
num_rows = table_out["predict_details"].get("num_rows", 0)
num_cols = table_out["predict_details"].get("num_cols", 0)
otsl_seq = (
table_out["predict_details"]
.get("prediction", {})
.get("rs_seq", [])
)
tbl = Table(
otsl_seq=otsl_seq,