diff --git a/docling/models/table_structure_model.py b/docling/models/table_structure_model.py index 4f471c4a..f17cbed0 100644 --- a/docling/models/table_structure_model.py +++ b/docling/models/table_structure_model.py @@ -209,23 +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 = 0 - num_cols = 0 - otsl_seq = [] - if "predict_details" in table_out: - if "num_rows" in table_out["predict_details"]: - num_rows = table_out["predict_details"]["num_rows"] - if "num_cols" in table_out["predict_details"]: - num_cols = table_out["predict_details"]["num_cols"] - if "prediction" in table_out["predict_details"]: - if ( - "rs_seq" - in table_out["predict_details"]["prediction"] - ): - 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,