diff --git a/dan/predict/attention.py b/dan/predict/attention.py index d5d5d747becf696a208b7ce79d2b9654b8deda37..117adb78cf0e76b6af85837d855f3ee420c9df50 100644 --- a/dan/predict/attention.py +++ b/dan/predict/attention.py @@ -190,7 +190,7 @@ def polygon_to_bbx(polygon): return [[x, y], [x + w, y], [x + w, y + h], [x, y + h]] -def threshold(mask, method="threshold", thresh=0): +def threshold(mask, method="otsu", thresh=0): """ Threshold a grayscale mask. :param mask: a grayscale image (np.array) @@ -208,7 +208,7 @@ def threshold(mask, method="threshold", thresh=0): # Blur and apply Otsu thresholding blur = cv2.GaussianBlur(mask, (15, 15), 0) _, bin_mask = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) - # Apply dilation with + # Apply dilation kernel_width = cv2.getStructuringElement( cv2.MORPH_CROSS, (max_kernel, min_kernel) ) @@ -239,6 +239,11 @@ def get_polygon(text, max_value, offset, weights, size=None, return_contours=Fal # Detect the objects contours contours, _ = cv2.findContours(bin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + if return_contours: + return {}, None + return {} + # Select best contour metrics = [compute_contour_metrics(coverage_vector, cnt) for cnt in contours] confidences, scores = map(list, zip(*metrics)) @@ -314,7 +319,9 @@ def plot_attention( (width, height), return_contours=True, ) - cv2.drawContours(coverage_vector, [contour], 0, (255), 3) + + if contour is not None: + cv2.drawContours(coverage_vector, [contour], 0, (255), 5) # Keep track of text length tot_len += len(text_piece) + offset