diff --git a/dan/ocr/predict/attention.py b/dan/ocr/predict/attention.py index 399f72b6fd94139ca44ab5eb9906ed4b5ffe5c29..6b78c81d3dc3a57b9cc0580536334442292ea3fe 100644 --- a/dan/ocr/predict/attention.py +++ b/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