Skip to content
Snippets Groups Projects
Commit d1c26db0 authored by Solene Tarride's avatar Solene Tarride
Browse files

always return contour

parent 6aa5bd9a
No related branches found
No related tags found
1 merge request!76Add predicted objects to predict command
This commit is part of merge request !76. Comments created here will be created in the context of that merge request.
......@@ -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(
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment