| """Detector pre/post-processing shared by evaluation and inference.""" |
|
|
| import cv2 |
| import numpy as np |
|
|
| from ..geometry import expand_quad, quad_size |
| from ..models.detector import KERNEL_END_TRIM, KERNEL_HEIGHT_RATIO |
|
|
| DET_MAX_SIDE = 1920 |
| DET_MIN_SIDE = 960 |
|
|
|
|
| def detector_input(img_rgb, max_side=DET_MAX_SIDE, min_side=DET_MIN_SIDE): |
| """Resize so the long side is within [min_side, max_side], pad to a multiple of 32. |
| |
| Returns (float32 CHW in [0,1], scale) where model coords = original coords * scale. |
| """ |
| h, w = img_rgb.shape[:2] |
| long = max(h, w) |
| scale = min(max_side / long, max(1.0, min_side / long)) |
| nh, nw = max(32, round(h * scale)), max(32, round(w * scale)) |
| img = cv2.resize(img_rgb, (nw, nh), interpolation=cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR) |
| ph, pw = (nh + 31) // 32 * 32, (nw + 31) // 32 * 32 |
| pad = np.empty((ph, pw, 3), np.uint8) |
| pad[:] = np.median(img[::8, ::8].reshape(-1, 3), 0).astype(np.uint8) |
| pad[:nh, :nw] = img |
| return pad.transpose(2, 0, 1).astype(np.float32) / 255.0, scale |
|
|
|
|
| def _text_quad(rect): |
| """minAreaRect -> quad ordered TL, TR, BR, BL with the top edge as close to horizontal as possible.""" |
| pts = cv2.boxPoints(rect) |
| c = pts.mean(0) |
| pts = pts[np.argsort(np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0]))] |
| best, best_score = 0, -1e9 |
| for s in range(4): |
| v = pts[(s + 1) % 4] - pts[s] |
| score = v[0] / (np.linalg.norm(v) + 1e-6) |
| if score > best_score: |
| best, best_score = s, score |
| return np.roll(pts, -best, axis=0).astype(np.float32) |
|
|
|
|
| def kernels_to_quads(prob, bin_thresh=0.3, box_thresh=0.55, min_area=4): |
| """prob: (H, W) kernel probability map. Returns list of (quad, score) in map coordinates.""" |
| bitmap = (prob > bin_thresh).astype(np.uint8) |
| n, labels, stats, _ = cv2.connectedComponentsWithStats(bitmap, connectivity=4) |
| out = [] |
| for i in range(1, n): |
| x, y, w, h, area = stats[i] |
| if area < min_area: |
| continue |
| comp = labels[y:y + h, x:x + w] == i |
| score = float(prob[y:y + h, x:x + w][comp].mean()) |
| if score < box_thresh: |
| continue |
| ys, xs = np.nonzero(comp) |
| pts = np.stack([xs + x, ys + y], 1).astype(np.float32) |
| (cx, cy), (rw, rh), ang = cv2.minAreaRect(pts) |
| q = _text_quad(((cx, cy), (rw, rh), ang)) |
| kw, kh = quad_size(q) |
| full_h = kh / KERNEL_HEIGHT_RATIO |
| dy = (full_h - kh) / 2 |
| dx = min(KERNEL_END_TRIM * full_h, kw / 2) |
| out.append((expand_quad(q, dx, dy, dx, dy), score)) |
| return out |
|
|
|
|
| def quad_iou(a, b): |
| a, b = np.asarray(a, np.float32), np.asarray(b, np.float32) |
| inter, _ = cv2.intersectConvexConvex(a, b) |
| union = abs(cv2.contourArea(a)) + abs(cv2.contourArea(b)) - inter |
| return float(inter / union) if union > 0 else 0.0 |
|
|
|
|
| def match_quads(pred, gt, thresh=0.5): |
| """Greedy one-to-one matching by IoU. Returns (matches [(pi, gi, iou)], unmatched_pred, unmatched_gt).""" |
| pairs = [] |
| for i, p in enumerate(pred): |
| pb = cv2.boundingRect(np.asarray(p, np.float32)) |
| for j, g in enumerate(gt): |
| gb = cv2.boundingRect(np.asarray(g, np.float32)) |
| if pb[0] > gb[0] + gb[2] or gb[0] > pb[0] + pb[2] or pb[1] > gb[1] + gb[3] or gb[1] > pb[1] + pb[3]: |
| continue |
| iou = quad_iou(p, g) |
| if iou >= thresh: |
| pairs.append((iou, i, j)) |
| pairs.sort(reverse=True) |
| used_p, used_g, matches = set(), set(), [] |
| for iou, i, j in pairs: |
| if i not in used_p and j not in used_g: |
| used_p.add(i) |
| used_g.add(j) |
| matches.append((i, j, iou)) |
| return matches, [i for i in range(len(pred)) if i not in used_p], [j for j in range(len(gt)) if j not in used_g] |
|
|
|
|
| def find_rules(img, text_height): |
| """Printed horizontal rules (thin, long dark lines) as [x0, y0, x1, y1] boxes. |
| |
| Newspapers and magazines separate stories with such rules; reading order uses them as separators. |
| Text survives the long horizontal opening only as underlines, which are harmless separators. |
| """ |
| gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) if img.ndim == 3 else img |
| h, w = gray.shape |
| ink = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 31, 12) |
| length = max(30, int(max(8 * text_height, w / 25))) |
| lines = _long_runs(ink, length) |
| lines = cv2.dilate(lines, cv2.getStructuringElement(cv2.MORPH_RECT, (max(3, int(text_height)), 1))) |
| n, _, stats, _ = cv2.connectedComponentsWithStats(lines, connectivity=8) |
| out = [] |
| for x, y, bw, bh, _ in stats[1:n]: |
| if bw >= length and bh <= max(4, 0.5 * text_height): |
| out.append([float(x), float(y), float(x + bw), float(y + bh)]) |
| return out |
|
|
|
|
| def _long_runs(ink, length): |
| """Binary opening with a 1 x length horizontal line, computed from run lengths: keeps exactly the horizontal |
| ink runs at least `length` long (several times faster than the sliding-window morphology on full pages).""" |
| b = np.zeros((ink.shape[0], ink.shape[1] + 2), np.int8) |
| b[:, 1:-1] = ink > 0 |
| d = np.diff(b, axis=1) |
| ys, x0 = np.nonzero(d == 1) |
| _, x1 = np.nonzero(d == -1) |
| keep = np.nonzero(x1 - x0 >= length)[0] |
| out = np.zeros_like(ink) |
| for k in keep: |
| out[ys[k], x0[k]:x1[k]] = 255 |
| return out |
|
|
|
|
| def split_gutter_crossers(gray, dets): |
| """Split detected lines that run across a column gutter (two columns' lines merged by the detector). |
| |
| A line is split only where (a) prose lines just above and below it leave a clear vertical gap - a |
| column gutter - on both sides, (b) the line itself has whitespace there, and (c) it is body-text sized |
| (headlines legitimately span columns). dets: list of (quad, score); returns the same. |
| """ |
| if len(dets) < 6: |
| return dets |
| boxes = np.array([[q[:, 0].min(), q[:, 1].min(), q[:, 0].max(), q[:, 1].max()] for q, _ in dets], np.float64) |
| hs = boxes[:, 3] - boxes[:, 1] |
| hmed = float(np.median(hs)) |
| cy = (boxes[:, 1] + boxes[:, 3]) / 2 |
| out = [] |
| for k, (q, score) in enumerate(dets): |
| x0, y0, x1, y1 = boxes[k] |
| if not (0.7 * hmed < hs[k] < 1.4 * hmed and x1 - x0 > 6 * hmed): |
| out.append((q, score)) |
| continue |
| near = (np.abs(cy - cy[k]) < 6 * hmed) & (boxes[:, 2] > x0) & (boxes[:, 0] < x1) |
| near[k] = False |
| if near.sum() < 4: |
| out.append((q, score)) |
| continue |
| lo = int(np.floor(x0)) |
| cov = np.zeros(int(np.ceil(x1)) - lo + 1, np.int32) |
| for bx0, _, bx1, _ in boxes[near]: |
| cov[max(0, int(bx0) - lo):max(0, int(bx1) - lo)] += 1 |
| split_at = None |
| free = np.flatnonzero(np.diff(np.r_[1, (cov == 0).astype(np.int8), 1])) |
| for a, b in zip(free[0::2], free[1::2]): |
| if b - a < 0.3 * hmed or a < hmed or b > len(cov) - hmed: |
| continue |
| gx0, gx1 = lo + a, lo + b |
| widths = boxes[:, 2] - boxes[:, 0] |
| ok = True |
| for side in (near & (boxes[:, 2] <= gx0 + 1), near & (boxes[:, 0] >= gx1 - 1)): |
| |
| |
| above, below = side & (cy < cy[k]), side & (cy > cy[k]) |
| if above.sum() < 1 or below.sum() < 1 or np.median(widths[side]) < 6 * hmed: |
| ok = False |
| if not ok: |
| continue |
| crop = gray[max(0, int(y0)):int(np.ceil(y1)), max(0, int(gx0 - 0.3 * hmed)):int(np.ceil(gx1 + 0.3 * hmed))] |
| if crop.size == 0: |
| continue |
| thr = min(200, int(np.median(crop)) - 40) |
| empty = (crop < thr).sum(0) == 0 |
| runs = np.flatnonzero(np.diff(np.r_[0, empty.astype(np.int8), 0])) |
| best = max(((e - s_, s_) for s_, e in zip(runs[0::2], runs[1::2])), default=(0, 0)) |
| if best[0] >= 0.25 * hmed: |
| split_at = int(gx0 - 0.3 * hmed) + best[1] + best[0] / 2 |
| break |
| if split_at is None: |
| out.append((q, score)) |
| continue |
| for a, b in ((x0, split_at - 1), (split_at + 1, x1)): |
| out.append((np.array([[a, y0], [b, y0], [b, y1], [a, y1]], np.float32), score)) |
| return out |
|
|
