| """Quadrilateral helpers. Quads are (4, 2) float arrays ordered TL, TR, BR, BL in text reading direction.""" |
|
|
| import cv2 |
| import numpy as np |
|
|
|
|
| def order_quad(pts) -> np.ndarray: |
| """Order 4 points TL, TR, BR, BL assuming text runs roughly left-to-right.""" |
| pts = np.asarray(pts, np.float32).reshape(4, 2) |
| s = pts.sum(1) |
| d = pts[:, 1] - pts[:, 0] |
| tl, br = pts[np.argmin(s)], pts[np.argmax(s)] |
| tr, bl = pts[np.argmin(d)], pts[np.argmax(d)] |
| q = np.stack([tl, tr, br, bl]) |
| if len({tuple(p) for p in q.tolist()}) < 4: |
| c = pts.mean(0) |
| ang = np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0]) |
| q = pts[np.argsort(ang)] |
| q = np.roll(q, -int(np.argmin(q.sum(1))), axis=0) |
| return q |
|
|
|
|
| def rect_quad(x0, y0, x1, y1) -> np.ndarray: |
| return np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], np.float32) |
|
|
|
|
| def quad_bbox(q) -> list[float]: |
| q = np.asarray(q) |
| return [float(q[:, 0].min()), float(q[:, 1].min()), float(q[:, 0].max()), float(q[:, 1].max())] |
|
|
|
|
| def quad_size(q): |
| """(width along text direction, height across it).""" |
| q = np.asarray(q, np.float32) |
| w = (np.linalg.norm(q[1] - q[0]) + np.linalg.norm(q[2] - q[3])) / 2 |
| h = (np.linalg.norm(q[3] - q[0]) + np.linalg.norm(q[2] - q[1])) / 2 |
| return float(w), float(h) |
|
|
|
|
| def expand_quad(q, left, top, right, bottom) -> np.ndarray: |
| """Move each side of the quad outward by the given amounts (negative shrinks), in pixels.""" |
| q = np.asarray(q, np.float32) |
| u = (q[1] - q[0]) + (q[2] - q[3]) |
| v = (q[3] - q[0]) + (q[2] - q[1]) |
| u /= max(np.linalg.norm(u), 1e-6) |
| v /= max(np.linalg.norm(v), 1e-6) |
| return np.stack([q[0] - u * left - v * top, q[1] + u * right - v * top, |
| q[2] + u * right + v * bottom, q[3] - u * left + v * bottom]).astype(np.float32) |
|
|
|
|
| def crop_quad(image, q, out_h=None) -> np.ndarray: |
| """Perspective-crop a quad to an upright rectangle. If out_h is given, scale to that height.""" |
| w, h = quad_size(q) |
| w, h = max(int(round(w)), 1), max(int(round(h)), 1) |
| if out_h is not None: |
| w, h = max(int(round(w * out_h / h)), 1), out_h |
| dst = rect_quad(0, 0, w, h) |
| m = cv2.getPerspectiveTransform(np.asarray(q, np.float32), dst) |
| return cv2.warpPerspective(image, m, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE) |
|
|
|
|
| def transform_quads(quads, m) -> list[np.ndarray]: |
| """Apply a 3x3 homography (or 2x3 affine) to each quad.""" |
| if not len(quads): |
| return [] |
| m = np.asarray(m, np.float64) |
| if m.shape == (2, 3): |
| m = np.vstack([m, [0, 0, 1]]) |
| if isinstance(quads, np.ndarray): |
| pts = quads.reshape(1, -1, 2).astype(np.float64) |
| else: |
| pts = np.concatenate([np.asarray(q, np.float64).reshape(4, 2) for q in quads])[None] |
| out = cv2.perspectiveTransform(pts, m)[0].astype(np.float32) |
| return [out[i * 4:(i + 1) * 4] for i in range(len(quads))] |
|
|
|
|
| def merge_line_cells(cells): |
| """cells: list of (x0, y0, x1, y1, text) in one layout box -> merged line segments.""" |
| |
| |
| rows = [] |
| for c in sorted(cells, key=lambda c: (c[1] + c[3]) / 2): |
| h = c[3] - c[1] |
| best, best_ov = None, 0.5 |
| for r in rows: |
| ov = (min(r["y1"], c[3]) - max(r["y0"], c[1])) / max(min(h, r["y1"] - r["y0"]), 1e-3) |
| if ov > best_ov: |
| best, best_ov = r, ov |
| if best is None: |
| rows.append({"y0": c[1], "y1": c[3], "cells": [c]}) |
| else: |
| best["cells"].append(c) |
| best["y0"], best["y1"] = min(best["y0"], c[1]), max(best["y1"], c[3]) |
| out = [] |
| for r in rows: |
| ln = sorted(r["cells"], key=lambda c: c[0]) |
| h = float(np.median([c[3] - c[1] for c in ln])) |
| groups = [[ln[0]]] |
| for prev, c in zip(ln, ln[1:]): |
| if c[0] - groups[-1][-1][2] > 0.6 * h: |
| groups.append([c]) |
| else: |
| groups[-1].append(c) |
| for g in groups: |
| text = g[0][4] |
| for prev, c in zip(g, g[1:]): |
| sep = " " if c[0] - prev[2] > 0.15 * h and not text.endswith(" ") and not c[4].startswith(" ") else "" |
| text += sep + c[4] |
| out.append((min(c[0] for c in g), min(c[1] for c in g), max(c[2] for c in g), max(c[3] for c in g), text)) |
| return out |
|
|
