From d1c26db005935a447bea98abe1e64bdf0bf9dd4d Mon Sep 17 00:00:00 2001
From: starride-teklia <starride@teklia.com>
Date: Tue, 11 Apr 2023 13:14:52 +0000
Subject: [PATCH] always return contour
---
dan/predict/attention.py | 14 ++++----------
1 file changed, 4 insertions(+), 10 deletions(-)
diff --git a/dan/predict/attention.py b/dan/predict/attention.py
index 4ef781c3..d1386694 100644
--- a/dan/predict/attention.py
+++ b/dan/predict/attention.py
@@ -112,7 +112,7 @@ def get_predicted_polygons_with_confidence(
start_index = 0
for text_piece, confidence in zip(text_list, confidence_list):
start_index += len(text_piece) + offset
- polygon = get_polygon(
+ polygon, _ = get_polygon(
text_piece, max_value, offset, weights, size=(width, height)
)
polygon["text"] = text_piece
@@ -220,14 +220,13 @@ def threshold(mask, method="otsu", thresh=0):
return bin_mask
-def get_polygon(text, max_value, offset, weights, size=None, return_contours=False):
+def get_polygon(text, max_value, offset, weights, size=None):
"""
Gets polygon associated with element of current text_piece, indexed by offset
:param text: Text piece selected with offset after splitting DAN prediction
:param max_value: Maximum "attention intensity" for parts of a text piece, used for normalization
:param offset: Offset value to get the relevant part of text piece
:param size: Target size (width, height) to resize the coverage vector
- :param return_contours: Return the contour of the current polygon (used for plotting)
"""
# Compute coverage vector
coverage_vector = compute_coverage(text, max_value, offset, weights, size=size)
@@ -239,9 +238,7 @@ def get_polygon(text, max_value, offset, weights, size=None, return_contours=Fal
contours, _ = cv2.findContours(bin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
- if return_contours:
- return {}, None
- return {}
+ return {}, None
# Select best contour
metrics = [compute_contour_metrics(coverage_vector, cnt) for cnt in contours]
@@ -257,10 +254,7 @@ def get_polygon(text, max_value, offset, weights, size=None, return_contours=Fal
}
simplified_contour = np.expand_dims(np.array(coord, dtype=np.int32), axis=1)
- if return_contours:
- return polygon, simplified_contour
-
- return polygon
+ return polygon, simplified_contour
def plot_attention(
--
GitLab