Ignore invalid objects during inference

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