NeoAnd/Docling
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Christoph Auer <cau@zurich.ibm.com>
Co-authored-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
Christoph Auer 2024-08-14 11:32:30 +02:00 committed by GitHub
parent dd0df9f094
commit 61be78a875
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
docling/models/layout_model.py
View File

@ -69,6 +69,10 @@ class LayoutModel:
"Key-Value Region": 0.45, "Key-Value Region": 0.45,
} }
CLASS_REMAPPINGS = {
"Document Index": "Table",
}
_log.debug("================= Start postprocess function ====================") _log.debug("================= Start postprocess function ====================")
start_time = time.time() start_time = time.time()
# Apply Confidence Threshold to cluster predictions # Apply Confidence Threshold to cluster predictions
@ -79,6 +83,10 @@ class LayoutModel:
confidence = CLASS_THRESHOLDS[cluster.label] confidence = CLASS_THRESHOLDS[cluster.label]
if cluster.confidence >= confidence: if cluster.confidence >= confidence:
# annotation["created_by"] = "high_conf_pred" # 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) clusters_out.append(cluster)
# map to dictionary clusters and cells, with bottom left origin # map to dictionary clusters and cells, with bottom left origin