Fix class re-mapping for table of contents (#33)
Signed-off-by: Christoph Auer <cau@zurich.ibm.com> Co-authored-by: Christoph Auer <cau@zurich.ibm.com>
This commit is contained in:
parent
dd0df9f094
commit
61be78a875
@ -69,6 +69,10 @@ class LayoutModel:
|
||||
"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 @@ class LayoutModel:
|
||||
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
|
||||
|
Loading…
Reference in New Issue
Block a user