Add predicted objects to predict command

Closes #36 (closed)

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