diff --git a/docling/models/layout_model.py b/docling/models/layout_model.py index 1b1f2246..93f80d54 100644 --- a/docling/models/layout_model.py +++ b/docling/models/layout_model.py @@ -69,6 +69,10 @@ def postprocess(self, clusters: List[Cluster], cells: List[Cell], page_height): "Key-Value Region": 0.45, } + CLASS_REMAPPINGS = { + "Document Index": "Table", + } + _log.debug("================= Start postprocess function ====================") start_time = time.time() # Apply Confidence Threshold to cluster predictions @@ -79,6 +83,10 @@ def postprocess(self, clusters: List[Cluster], cells: List[Cell], page_height): confidence = CLASS_THRESHOLDS[cluster.label] if cluster.confidence >= confidence: # annotation["created_by"] = "high_conf_pred" + + # Remap class labels where needed. + if cluster.label in CLASS_REMAPPINGS.keys(): + cluster.label = CLASS_REMAPPINGS[cluster.label] clusters_out.append(cluster) # map to dictionary clusters and cells, with bottom left origin