From 6f00921df245ed5222ddefaac5e591eeb0b1b222 Mon Sep 17 00:00:00 2001
From: manonBlanco <blanco@teklia.com>
Date: Tue, 19 Dec 2023 11:49:00 +0100
Subject: [PATCH] Ignore invalid objects during inference
---
dan/ocr/predict/attention.py | 12 ++++++++++--
1 file changed, 10 insertions(+), 2 deletions(-)
diff --git a/dan/ocr/predict/attention.py b/dan/ocr/predict/attention.py
index 399f72b6..6b78c81d 100644
--- a/dan/ocr/predict/attention.py
+++ b/dan/ocr/predict/attention.py
@@ -253,6 +253,8 @@ def get_predicted_polygons_with_confidence(
size=(width, height),
)
start_index += len(text_piece) + offset
+ if not polygon:
+ continue
polygon["text"] = text_piece
polygon["text_confidence"] = confidence
polygons.append(polygon)
@@ -362,7 +364,7 @@ def get_polygon(
weights: np.ndarray,
size: Tuple[int, int] | None = None,
max_object_height: int = 50,
-) -> Tuple[dict, np.ndarray]:
+) -> Tuple[dict, np.ndarray | None]:
"""
Gets polygon associated with element of current text_piece, indexed by offset
:param text: Text piece selected with offset after splitting DAN prediction
@@ -382,6 +384,8 @@ def get_polygon(
if max_object_height
else get_best_contour(coverage_vector, bin_mask)
)
+ if not coord or confidence is None:
+ return {}, None
# Format for JSON
polygon = {
@@ -400,7 +404,7 @@ def get_best_contour(coverage_vector, bin_mask):
contours, _ = cv2.findContours(bin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
- return {}, None
+ return [], None
# Select best contour
metrics = [compute_contour_metrics(coverage_vector, cnt) for cnt in contours]
@@ -417,6 +421,10 @@ def get_grid_search_contour(coverage_vector, bin_mask, height=50):
"""
# Limit search area based on attention values
roi = np.argwhere(bin_mask == 255)
+
+ if not np.any(roi):
+ return [], None
+
y_min, y_max = roi[:, 0].min(), roi[:, 0].max()
# Limit bounding box shape
--
GitLab