fix: fixed boxes sorting error

This commit is contained in:
SWHL committed 2025-12-02 11:26:16 +08:00
1 parent 3f3fb13235
commit 7c1e4601b8
2 files changed
+17 -27

No files matched your search

+2 -10
View File
@@ -1,17 +1,9 @@
# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
from rapidocr import EngineType, RapidOCR
from rapidocr import RapidOCR
engine = RapidOCR(
params={
"Det.engine_type": EngineType.TORCH,
"Cls.engine_type": EngineType.TORCH,
"Rec.engine_type": EngineType.TORCH,
"EngineConfig.torch.use_cuda": True, # 使用torch GPU版推理
"EngineConfig.torch.gpu_id": 0, # 指定GPU id
}
)
engine = RapidOCR()
img_url = "https://github.com/RapidAI/RapidOCR/blob/main/python/tests/test_files/ch_en_num.jpg?raw=true"
result = engine(img_url)
+15 -17
View File
@@ -24,7 +24,6 @@ from rapidocr.inference_engine.base import get_engine
from .utils import DBPostProcess, DetPreProcess, TextDetOutput
_BOX_SORT_Y_THRESHOLD = 10
_BOX_SORT_LINE_SEPARATION_FACTOR = 1e6
class TextDetector:
@@ -82,25 +81,24 @@ class TextDetector:
@staticmethod
def sorted_boxes(dt_boxes: np.ndarray) -> np.ndarray:
"""
Sort text boxes in order from top to bottom, left to right
args:
dt_boxes(array):detected text boxes with shape [4, 2]
return:
sorted boxes(array) with shape [4, 2]
Equivalent NumPy implementation of the original bubble-adjusted sort.
"""
if len(dt_boxes) == 0:
return dt_boxes
# Sort by y, then identify lines, then sort by (line, x)
y_order = np.argsort(dt_boxes[:, 0, 1], kind="stable")
sorted_y = dt_boxes[y_order, 0, 1]
# Step 1: Stable sort by y (top to bottom)
y_coords = dt_boxes[:, 0, 1]
y_order = np.argsort(y_coords, kind="stable")
boxes_y_sorted = dt_boxes[y_order]
y_sorted = y_coords[y_order]
line_ids = np.empty(len(dt_boxes), dtype=np.int32)
line_ids[0] = 0
np.cumsum(np.abs(np.diff(sorted_y)) >= _BOX_SORT_Y_THRESHOLD, out=line_ids[1:])
# Step 2: Assign line IDs based on adjacent y differences
dy = np.diff(y_sorted)
line_increments = (dy >= _BOX_SORT_Y_THRESHOLD).astype(np.int32)
line_ids = np.concatenate([[0], np.cumsum(line_increments)])
# Create composite sort key for final ordering
# Shift line_ids by large factor, add x for tie-breaking
sort_key = line_ids[y_order] * _BOX_SORT_LINE_SEPARATION_FACTOR + dt_boxes[y_order, 0, 0]
final_order = np.argsort(sort_key, kind="stable")
return dt_boxes[y_order[final_order]]
# Now, within each line_id group, sort by x (left to right)
x_coords = boxes_y_sorted[:, 0, 0]
final_order_in_y_sorted = np.lexsort((x_coords, line_ids))
return boxes_y_sorted[final_order_in_y_sorted]