Fix class re-mapping for table of contents (#33)

Signed-off-by: Christoph Auer <[email protected]>
Co-authored-by: Christoph Auer <[email protected]>

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