Download code/hlt_postproc_quads.py from tailot/hlt-geometric-vision: direct link, hf CLI and curl.
- Browser
- Download file 35 kB
-
https://huggingface.co/tailot/hlt-geometric-vision/resolve/main/code/hlt_postproc_quads.py
- Command line
-
hf download hf://tailot/hlt-geometric-vision/code/hlt_postproc_quads.py
-
curl -L -o hlt_postproc_quads.py https://huggingface.co/tailot/hlt-geometric-vision/resolve/main/code/hlt_postproc_quads.py
35 kB
| """Extraction of ALL quadrilaterals (rectangles, even in perspective) from a real edge map. | |
| Why it exists: on the real domain the UNet edge_prob transfers, the keypoint_prob does not, and | |
| the 1-polygon pipeline (max_polygons=1) merges the corners of different objects into a single | |
| cloud (star-shaped polygon). Here quadrilaterals are rebuilt from the edge map alone: | |
| 1. dominant lines: Hough weighted by edge_prob + NMS, then REFIT to the edge pixels (TLS); | |
| 2. candidates = ANY 4 lines in cyclic order (a,b,c,d); the corners are the 4 consecutive | |
| intersections (sub-pixel). VECTORIZED enumeration (C(N,4) x 3 orders); | |
| 3. score = fraction of the side covered by edge, measured in O(1) with cumulative sums along | |
| each line, with a threshold on the worst side too; | |
| 4. shape: "rect" = interior angles in [50,130] degrees (rectangles even in perspective), | |
| "quad" = any convex quadrilateral (trapezoids, parallelograms, rhombi); | |
| 5. contrast test with coherent polarity (discards rectangles "stitched" together from | |
| uncorrelated edges). Sides on the image border are NOT penalized: off canvas there is no pixel; | |
| 6. NMS: by default the OUTER contour is kept (the shapes inside its content are | |
| discarded); with nested=True nested rectangles are kept too (inner window of a | |
| picture frame); | |
| 7. refinement: each side refitted to the edge pixels and the corners recomputed. | |
| Differences w.r.t. the previous version (measured causes of missed detections): | |
| - max_quads=8 truncated the 12 frames of a grid down to 8; | |
| - the 40 best lines were saturated by background horizontals (29 hor / 9 ver on a.jpg): the | |
| verticals of the paintings did not enter the pool; | |
| - the contrast test rejected sides flush with the image border (cropped object): off canvas | |
| the "outer" sample was the edge pixel itself => contrast 0; | |
| - the constant gray pad created two horizontal lines as long as the image. | |
| Coordinates in pixels of the image passed in (x to the right, y downward). Corner order: | |
| clockwise, starting from the corner with the smallest x+y. | |
| """ | |
| from __future__ import annotations | |
| import itertools | |
| from dataclasses import dataclass | |
| from typing import Dict, List, Optional, Tuple | |
| import cv2 | |
| import numpy as np | |
| class QuadConfig: | |
| edge_thresh: float = 0.5 # edge pixels = edge_prob > threshold | |
| n_theta: int = 360 # angular bins over [0, pi): 0.5 degrees | |
| rho_step: float = 1.0 | |
| n_lines: int = 500 # lines kept after NMS (grows with valid segments, not C(N,4)) | |
| nms_theta_bins: int = 6 # suppression +/-3 degrees | |
| nms_rho_bins: int = 6 # suppression +/-6 px | |
| min_line_votes: float = 25.0 # minimum sum of edge_prob along the line | |
| min_rel_votes: float = 0.05 # and at least 5% of the strongest line | |
| refit_band_px: float = 3.0 # the line refit uses edge pixels within +/-3 px | |
| shape: str = "rect" # "rect" | "quad" | |
| corner_angle_deg: Tuple[float, float] = (50.0, 130.0) # only shape="rect" | |
| band_px: float = 3.0 # spatial tolerance of the support | |
| trim: float = 0.06 # 6% is ignored at each end of the side (bitten corners) | |
| min_side_px: float = 30.0 | |
| min_area_frac: float = 0.01 | |
| max_area_frac: float = 0.60 | |
| max_aspect: float = 6.0 | |
| border_margin_px: float = 8.0 # tolerance for corners off canvas (cropped objects) | |
| min_side_cov: float = 0.25 # minimum edge coverage of the WORST side (see hybrid()) | |
| min_mean_cov: float = 0.50 # minimum mean edge coverage of the perimeter | |
| contrast_px: Tuple[int, int] = (3, 8) # inside/outside band of the side where gray is measured | |
| contrast_min: float = 12.0 # minimum inside-outside contrast in gray levels | |
| min_contrast_cov: float = 0.0 # 0 = polarity test DISABLED (fragile on paintings: see README) | |
| vote_halfwidth_deg: float = 6.0 # DIRECT Hough: each pixel votes only for theta within +/- this | |
| vote_coherence: float = 0.5 # a pixel votes only if its local orientation is coherent enough | |
| align_min: float = 0.6 # minimum edge orientation coherence w.r.t. the line | |
| align_win: int = 7 # window (px) over which edge orientation is averaged | |
| max_gap_frac: float = 0.20 # largest gap without aligned edge, as a fraction of side length | |
| prefer_outer: bool = True # among near-duplicates keeps the outer contour (frame border) | |
| prefer_outer_tol: float = 0.15 # score tolerance for prefer_outer | |
| stretch_margin: float = 0.02 # extends a child (>=2 collinear sides) at lower score: artifact | |
| group_cover: float = 0.60 # 'stitched group' if >=2 disjoint children cover >=60% of it | |
| max_eval: int = 6000 # how many candidates (best by coverage) reach the slow checks | |
| nested: bool = False # True: also keep nested rectangles (e.g. inner window) | |
| nms_iou: float = 0.30 | |
| nms_contain: float = 0.70 # candidate >=70% contained in an already chosen one -> dropped | |
| max_quads: int = 64 | |
| refine_band_px: float = 4.0 | |
| refine_max_shift_px: float = 10.0 | |
| chunk: int = 80000 # line combinations per block (memory) | |
| def hybrid(cls, **kw) -> "QuadConfig": | |
| """Thresholds for the U-net + Canny map: the missing sides are filled in by Canny, so the | |
| coverage of true quads rises to >=0.93 and much more can be demanded. | |
| CALIBRATED ON A SMALL SET (see README): they must be retuned on a set with GT.""" | |
| base = dict(min_side_cov=0.60, min_mean_cov=0.75) | |
| base.update(kw) | |
| return cls(**base) | |
| # ------------------------------------------------------------------ lines | |
| def hough_lines(edge: np.ndarray, cfg: QuadConfig, orient=None) -> List[Tuple[float, float, float]]: | |
| """Lines (rho, theta, votes) from Hough weighted by edge_prob. x cos t + y sin t = rho. | |
| With orient=(C,S,M) the vote is DIRECT: each edge pixel votes only for the theta within | |
| +/- vote_halfwidth_deg of its own local normal, and only if the local orientation is | |
| coherent. Texture curves (brush strokes, grain) then do not accumulate votes like a | |
| straight line and do not saturate the pool (measured: on a.jpg 200 lines were not enough).""" | |
| h, w = edge.shape | |
| ys, xs = np.nonzero(edge > cfg.edge_thresh) | |
| if xs.size == 0: | |
| return [] | |
| wts = edge[ys, xs].astype(np.float64) | |
| dth = np.pi / cfg.n_theta | |
| nr = int(2 * np.ceil(np.hypot(h, w) / cfg.rho_step) + 1) | |
| r0 = (nr - 1) / 2.0 | |
| if orient is None: | |
| th = np.arange(cfg.n_theta) * dth | |
| rho = xs[:, None] * np.cos(th)[None] + ys[:, None] * np.sin(th)[None] | |
| ridx = np.clip(np.rint(rho / cfg.rho_step + r0).astype(np.int64), 0, nr - 1) | |
| tidx = np.broadcast_to(np.arange(cfg.n_theta)[None, :], ridx.shape) | |
| wv = np.repeat(wts, cfg.n_theta) | |
| else: | |
| C, S_, M = orient | |
| c, s_, m = C[ys, xs], S_[ys, xs], M[ys, xs] | |
| coh = np.hypot(c, s_) / np.maximum(m, 1e-6) | |
| keep = (coh >= cfg.vote_coherence) & (m > 1e-6) | |
| xs, ys, wts, c, s_ = xs[keep], ys[keep], wts[keep], c[keep], s_[keep] | |
| if xs.size == 0: | |
| return [] | |
| t0 = 0.5 * np.arctan2(s_, c) # local normal, mod pi | |
| half = int(round(np.radians(cfg.vote_halfwidth_deg) / dth)) | |
| ks = np.arange(-half, half + 1) | |
| tk = t0[:, None] + ks[None, :] * dth # (n, K) continuous theta | |
| # wrap theta into [0, pi): rho is computed AFTER, with the wrapped theta (cos/sin flip | |
| # their sign on their own: no further sign inversion here) | |
| tk = np.where(tk < 0, tk + np.pi, tk) | |
| tk = np.where(tk >= np.pi, tk - np.pi, tk) | |
| rho = xs[:, None] * np.cos(tk) + ys[:, None] * np.sin(tk) | |
| tidx = np.rint(tk / dth).astype(np.int64) | |
| wrap = tidx >= cfg.n_theta # theta ~ pi == theta 0, rho flipped | |
| tidx = np.where(wrap, tidx - cfg.n_theta, tidx) | |
| rho = np.where(wrap, -rho, rho) | |
| ridx = np.clip(np.rint(rho / cfg.rho_step + r0).astype(np.int64), 0, nr - 1) | |
| wv = np.repeat(wts * coh[keep], ks.size) | |
| flat = (tidx * nr + ridx).ravel() | |
| acc = np.bincount(flat, weights=wv, minlength=cfg.n_theta * nr) | |
| acc = acc.reshape(cfg.n_theta, nr).astype(np.float32) | |
| work = cv2.boxFilter(acc, -1, (3, 1), normalize=False) # SUM over 3 rho bins (not the mean) | |
| peak0 = float(work.max()) | |
| out: List[Tuple[float, float, float]] = [] | |
| for _ in range(cfg.n_lines * 3): # margin: duplicates are removed later | |
| ti, ri = np.unravel_index(int(work.argmax()), work.shape) | |
| v = float(work[ti, ri]) | |
| if v < max(cfg.min_line_votes, cfg.min_rel_votes * peak0): | |
| break | |
| t0_, t1_ = max(0, ti - 2), min(cfg.n_theta, ti + 3) | |
| r_lo, r_hi = max(0, ri - 2), min(nr, ri + 3) | |
| win = np.clip(work[t0_:t1_, r_lo:r_hi], 0, None) ** 2 | |
| sw = float(win.sum()) or 1.0 | |
| tt = float((win.sum(axis=1) * np.arange(t0_, t1_)).sum() / sw) | |
| rr = float((win.sum(axis=0) * np.arange(r_lo, r_hi)).sum() / sw) | |
| out.append(((rr - r0) * cfg.rho_step, tt * dth, v)) | |
| work[max(0, ti - cfg.nms_theta_bins):ti + cfg.nms_theta_bins + 1, | |
| max(0, ri - cfg.nms_rho_bins):ri + cfg.nms_rho_bins + 1] = 0.0 | |
| return _dedupe(out)[: cfg.n_lines] | |
| def _dedupe(lines: List[Tuple[float, float, float]]) -> List[Tuple[float, float, float]]: | |
| """Removes duplicates, including theta~0 vs theta~pi (same line, opposite rho sign).""" | |
| if not lines: | |
| return [] | |
| L = sorted(lines, key=lambda z: -z[2]) | |
| r = np.array([z[0] for z in L]); t = np.array([z[1] for z in L]) | |
| px, py = r * np.cos(t), r * np.sin(t) # point on the line | |
| dth = np.abs(((t[:, None] - t[None, :]) + np.pi / 2) % np.pi - np.pi / 2) | |
| dist = np.abs(px[:, None] * np.cos(t)[None, :] + py[:, None] * np.sin(t)[None, :] - r[None, :]) | |
| dup = (dth < np.radians(2.5)) & (dist < 6.0) | |
| dead = np.zeros(len(L), bool) | |
| keep = [] | |
| for i in range(len(L)): | |
| if dead[i]: | |
| continue | |
| keep.append(L[i]) | |
| dead[i + 1:] |= dup[i, i + 1:] | |
| return keep | |
| def _fit_line(pts: np.ndarray, w: np.ndarray): | |
| """Weighted total least squares line. Returns (rho, theta) with theta in [0, pi).""" | |
| c = (pts * w[:, None]).sum(0) / w.sum() | |
| d = pts - c | |
| cov = (d * w[:, None]).T @ d / w.sum() | |
| _, evec = np.linalg.eigh(cov) | |
| direction = evec[:, 1] # principal axis | |
| n = np.array([-direction[1], direction[0]]) | |
| theta = float(np.arctan2(n[1], n[0])) | |
| rho = float(c @ n) | |
| if theta < 0: | |
| theta += np.pi | |
| rho = -rho | |
| if theta >= np.pi: | |
| theta -= np.pi | |
| rho = -rho | |
| return rho, theta | |
| def _refit_lines(edge: np.ndarray, lines, cfg: QuadConfig, iters: int = 2): | |
| """Hough lines are quantized to 0.5 degrees / 1 px: over 500 px that is ~2 px of error, | |
| which is almost the whole support band. They are refitted to the nearby edge pixels (TLS).""" | |
| ys, xs = np.nonzero(edge > cfg.edge_thresh) | |
| if xs.size == 0: | |
| return lines | |
| pts = np.stack([xs, ys], 1).astype(np.float64) | |
| wts = edge[ys, xs].astype(np.float64) | |
| out = [] | |
| for r, t, v in lines: | |
| for _ in range(iters): | |
| n = np.array([np.cos(t), np.sin(t)]) | |
| m = np.abs(pts @ n - r) <= cfg.refit_band_px | |
| if m.sum() < 12: | |
| break | |
| r, t = _fit_line(pts[m], wts[m]) | |
| out.append((r, t, v)) | |
| return _dedupe(out) | |
| # ------------------------------------------------------- vector geometry | |
| def _line_geometry(lines, h: int, w: int): | |
| """For N lines: base point p0, direction u, intersection matrices (PX,PY) and parameter T[i,j] | |
| = curvilinear abscissa, along line i, of the intersection with line j (nan if parallel).""" | |
| L = np.asarray([[r, t] for r, t, _ in lines], np.float64) | |
| r, t = L[:, 0], L[:, 1] | |
| p0 = np.stack([r * np.cos(t), r * np.sin(t)], 1) | |
| u = np.stack([-np.sin(t), np.cos(t)], 1) | |
| det = np.sin(t[None, :] - t[:, None]) # det[i,j] = sin(tj - ti) | |
| with np.errstate(divide="ignore", invalid="ignore"): | |
| PX = (r[:, None] * np.sin(t)[None, :] - r[None, :] * np.sin(t)[:, None]) / det | |
| PY = (r[None, :] * np.cos(t)[:, None] - r[:, None] * np.cos(t)[None, :]) / det | |
| bad = np.abs(det) < 1e-3 # ~0.06 degrees: nearly parallel | |
| PX[bad] = np.nan | |
| PY[bad] = np.nan | |
| T = (PX - p0[:, 0:1]) * u[:, 0:1] + (PY - p0[:, 1:2]) * u[:, 1:2] | |
| return p0, u, PX, PY, T | |
| def _orientation_fields(edge: np.ndarray, cfg: QuadConfig): | |
| """Local orientation of the edges, as sums over the double angle (invariant to +/-gradient): | |
| C = sum |g| cos 2t, S = sum |g| sin 2t, M = sum |g| over an align_win x align_win window. | |
| Needed to count as 'support' of a line ONLY the edges laid out along it: the texture of a | |
| painting (brush strokes) has random orientation and must not cover any line.""" | |
| e = cv2.GaussianBlur(np.asarray(edge, np.float32), (0, 0), 1.5) | |
| gx = cv2.Sobel(e, cv2.CV_32F, 1, 0, ksize=3) | |
| gy = cv2.Sobel(e, cv2.CV_32F, 0, 1, ksize=3) | |
| mag = np.hypot(gx, gy) | |
| ang2 = 2.0 * np.arctan2(gy, gx) | |
| k = (cfg.align_win, cfg.align_win) | |
| C = cv2.boxFilter(mag * np.cos(ang2), -1, k, normalize=False) | |
| S = cv2.boxFilter(mag * np.sin(ang2), -1, k, normalize=False) | |
| M = cv2.boxFilter(mag, -1, k, normalize=False) | |
| return C, S, M | |
| def _profiles(near: np.ndarray, orient, lines, p0: np.ndarray, u: np.ndarray, S: int, | |
| cfg: QuadConfig): | |
| """prof[i, k] = 1 if sample k of line i (step 1 px, s in [-S, S]) falls near an edge whose | |
| orientation is coherent with the line (normal = theta). Returns (cum, prof): cum is the | |
| cumulative sum (coverage of a segment in O(1)); prof is used to measure the gaps.""" | |
| h, w = near.shape | |
| C, Sn, M = orient | |
| s = np.arange(-S, S + 1, dtype=np.float64) | |
| xs = p0[:, 0:1] + u[:, 0:1] * s[None] | |
| ys = p0[:, 1:2] + u[:, 1:2] * s[None] | |
| xi, yi = np.rint(xs).astype(np.int64), np.rint(ys).astype(np.int64) | |
| ok = (xi >= 0) & (xi < w) & (yi >= 0) & (yi < h) | |
| xc, yc = np.clip(xi, 0, w - 1), np.clip(yi, 0, h - 1) | |
| phi2 = 2.0 * np.asarray([t for _, t, _ in lines])[:, None] # normal = theta | |
| align = (C[yc, xc] * np.cos(phi2) + Sn[yc, xc] * np.sin(phi2)) / np.maximum(M[yc, xc], 1e-6) | |
| prof = (ok & near[yc, xc] & (align >= cfg.align_min)).astype(np.float32) | |
| cum = np.zeros((prof.shape[0], prof.shape[1] + 1), np.float32) | |
| np.cumsum(prof, axis=1, out=cum[:, 1:]) | |
| return cum, prof | |
| def _cov_vec(cum, S, L, s1, s2, trim): | |
| """Mean coverage of the segment [s1,s2] of line L (array indices), cutting trim off each | |
| end. Returns (coverage, i1, i2): the indices find the segment again inside the profile.""" | |
| lo, hi = np.minimum(s1, s2), np.maximum(s1, s2) | |
| span = hi - lo | |
| i1 = np.clip(np.rint(lo + trim * span).astype(np.int64) + S, 0, 2 * S) | |
| i2 = np.clip(np.rint(hi - trim * span).astype(np.int64) + S, 0, 2 * S) | |
| n = np.maximum(i2 - i1 + 1, 1) | |
| return (cum[L, i2 + 1] - cum[L, i1]) / n, i1, i2 | |
| def _longest_gap(prof_row: np.ndarray, i1: int, i2: int) -> int: | |
| """Length (px) of the longest consecutive stretch WITHOUT support in prof_row[i1..i2].""" | |
| seg = prof_row[i1:i2 + 1] > 0 | |
| idx = np.nonzero(np.concatenate(([True], seg, [True])))[0] | |
| return int(np.diff(idx).max() - 1) | |
| def _poly_area(c: np.ndarray) -> float: | |
| x, y = c[:, 0], c[:, 1] | |
| return 0.5 * abs(float(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))) | |
| def _overlap(a: np.ndarray, b: np.ndarray) -> Tuple[float, float]: | |
| """(IoU, fraction of a contained in b) for two convex quadrilaterals.""" | |
| inter, _ = cv2.intersectConvexConvex(a.astype(np.float32), b.astype(np.float32)) | |
| inter = float(inter) | |
| aa, ab = _poly_area(a), _poly_area(b) | |
| return inter / max(aa + ab - inter, 1e-6), inter / max(aa, 1e-6) | |
| def canonical_order(c: np.ndarray) -> np.ndarray: | |
| """Clockwise (in image coordinates, y downward), first corner = min(x+y).""" | |
| c = np.asarray(c, np.float64) | |
| x, y = c[:, 0], c[:, 1] | |
| if float(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1))) < 0: | |
| c = c[::-1] | |
| k = int(np.argmin(c[:, 0] + c[:, 1])) | |
| return np.roll(c, -k, axis=0) | |
| # ------------------------------------------------------- contrast test | |
| def _side_contrast(gray: np.ndarray, p: np.ndarray, q: np.ndarray, centroid: np.ndarray, | |
| cfg: QuadConfig) -> Tuple[np.ndarray, np.ndarray]: | |
| """For each sample along the side p->q: (mean gray INSIDE - mean gray OUTSIDE, valid). | |
| A sample is NOT valid if the inner band or the outer band leaves the image: there the | |
| contrast cannot be measured (cropped object) and must not count as 'absence of edge'.""" | |
| h, w = gray.shape | |
| seg = q - p | |
| length = float(np.hypot(*seg)) | |
| u = seg / max(length, 1e-9) | |
| nrm = np.array([-u[1], u[0]]) | |
| if float((centroid - p) @ nrm) < 0: | |
| nrm = -nrm # normal pointing to the inside | |
| n = max(4, int(length / 3.0)) | |
| t = np.linspace(cfg.trim, 1.0 - cfg.trim, n) | |
| pts = p[None] + t[:, None] * seg[None] | |
| offs = np.arange(cfg.contrast_px[0], cfg.contrast_px[1] + 1, dtype=np.float64) | |
| def side_mean(sign: float): | |
| xs = pts[:, 0:1] + sign * nrm[0] * offs[None] | |
| ys = pts[:, 1:2] + sign * nrm[1] * offs[None] | |
| xr, yr = np.rint(xs).astype(int), np.rint(ys).astype(int) | |
| inside = ((xr >= 0) & (xr < w) & (yr >= 0) & (yr < h)).all(axis=1) | |
| xi, yi = np.clip(xr, 0, w - 1), np.clip(yr, 0, h - 1) | |
| return gray[yi, xi].mean(axis=1), inside | |
| a, va = side_mean(+1.0) | |
| b, vb = side_mean(-1.0) | |
| return a - b, va & vb | |
| def _contrast_support(gray: np.ndarray, c: np.ndarray, cfg: QuadConfig) -> Tuple[List[float], int]: | |
| """Coverage per side of the contrast with polarity COHERENT over all 4 sides. A real | |
| object is (e.g.) darker inside on every side; a rectangle stitched from uncorrelated | |
| edges has mixed polarity. Samples that cannot be measured (off canvas) count as passed.""" | |
| cen = c.mean(axis=0) | |
| sides = [_side_contrast(gray, c[i], c[(i + 1) % 4], cen, cfg) for i in range(4)] | |
| best, best_pol = [0.0] * 4, 1 | |
| for pol in (1, -1): | |
| cc = [float((~v | (pol * d > cfg.contrast_min)).mean()) for d, v in sides] | |
| if min(cc) >= min(best) and (min(cc) > min(best) or sum(cc) > sum(best)): | |
| best, best_pol = cc, pol | |
| return best, best_pol | |
| # --------------------------------------------------------------- refinement | |
| def _intersect(l1, l2): | |
| (r1, t1), (r2, t2) = l1[:2], l2[:2] | |
| det = np.sin(t2 - t1) | |
| if abs(det) < 1e-6: | |
| return None | |
| return np.array([(r1 * np.sin(t2) - r2 * np.sin(t1)) / det, | |
| (r2 * np.cos(t1) - r1 * np.cos(t2)) / det]) | |
| def _refine(edge: np.ndarray, corners: np.ndarray, cfg: QuadConfig) -> Tuple[np.ndarray, float]: | |
| ys, xs = np.nonzero(edge > cfg.edge_thresh) | |
| pts = np.stack([xs, ys], 1).astype(np.float64) | |
| wts = edge[ys, xs].astype(np.float64) | |
| lines = [] | |
| for i in range(4): | |
| p, q = corners[i], corners[(i + 1) % 4] | |
| seg = q - p | |
| L = float(np.hypot(*seg)) | |
| u = seg / max(L, 1e-9) | |
| nvec = np.array([-u[1], u[0]]) | |
| rel = pts - p[None] | |
| s = rel @ u | |
| d = np.abs(rel @ nvec) | |
| m = (d <= cfg.refine_band_px) & (s >= 0.08 * L) & (s <= 0.92 * L) | |
| if m.sum() >= 8: | |
| lines.append(_fit_line(pts[m], wts[m])) | |
| else: | |
| c = (p + q) / 2 | |
| th = float(np.arctan2(nvec[1], nvec[0])); rh = float(c @ nvec) | |
| if th < 0: | |
| th += np.pi; rh = -rh | |
| lines.append((rh, th)) | |
| new = [] | |
| for i in range(4): | |
| pt = _intersect(lines[(i - 1) % 4], lines[i]) | |
| new.append(corners[i] if pt is None or np.hypot(*(pt - corners[i])) > cfg.refine_max_shift_px | |
| else pt) | |
| new = np.array(new) | |
| return new, float(np.linalg.norm(new - corners, axis=1).mean()) | |
| # ----------------------------------------------------------------- main | |
| def _valid_segments(lines, cum, S, T, PX, PY, h, w, cfg: QuadConfig): | |
| """For each line L: the segments [intersection with j, intersection with k] covered by an | |
| aligned edge >= min_side_cov and at least min_side_px long. A segment is a possible SIDE. | |
| Returns seg[L] = {(j, k) with j<k: coverage}. It grows with the number of VALID segments, | |
| not with C(N,4): so hundreds of lines can be kept, hence short sides of small objects too.""" | |
| N = len(lines) | |
| m = cfg.border_margin_px | |
| inside = (np.isfinite(PX) & np.isfinite(PY) & (PX >= -m) & (PX <= w - 1 + m) | |
| & (PY >= -m) & (PY <= h - 1 + m)) | |
| seg: List[Dict[Tuple[int, int], float]] = [dict() for _ in range(N)] | |
| for L in range(N): | |
| mj = np.nonzero(inside[L])[0] | |
| mj = mj[mj != L] | |
| if mj.size < 2: | |
| continue | |
| ii, kk = np.triu_indices(mj.size, 1) | |
| t1, t2 = T[L, mj[ii]], T[L, mj[kk]] | |
| length = np.abs(t2 - t1) | |
| ok = length >= cfg.min_side_px | |
| if not ok.any(): | |
| continue | |
| ii, kk, t1, t2 = ii[ok], kk[ok], t1[ok], t2[ok] | |
| cov, _, _ = _cov_vec(cum, S, np.full(ii.shape, L), t1, t2, cfg.trim) | |
| good = cov >= cfg.min_side_cov | |
| for a, b, c in zip(mj[ii[good]], mj[kk[good]], cov[good]): | |
| seg[L][(int(min(a, b)), int(max(a, b)))] = float(c) | |
| return seg | |
| def _enumerate(lines, near, orient, h, w, cfg: QuadConfig): | |
| """All convex quadrilaterals with 4 sides covered enough, as 4-cycles of valid segments: | |
| side on b (between a and c), on c (between b and d), on d (between c and a), on a (between | |
| d and b). Returns (list of (mean_cov, min_cov, corners(4,2), sides_cov(4), side_info), prof) | |
| with side_info = [(line, i1, i2)] x 4 to locate each side in the profile prof.""" | |
| N = len(lines) | |
| if N < 4: | |
| return [], None | |
| S = int(np.ceil(np.hypot(h, w))) | |
| p0, u, PX, PY, T = _line_geometry(lines, h, w) | |
| cum, prof = _profiles(near, orient, lines, p0, u, S, cfg) | |
| seg = _valid_segments(lines, cum, S, T, PX, PY, h, w, cfg) | |
| # adj[c][b] = {d : segment on c between b and d is valid} | |
| adj: List[Dict[int, List[int]]] = [dict() for _ in range(N)] | |
| for L in range(N): | |
| for (j, k) in seg[L]: | |
| adj[L].setdefault(j, []).append(k) | |
| adj[L].setdefault(k, []).append(j) | |
| cyc = [] | |
| for b in range(N): | |
| for (a, c) in seg[b]: | |
| if a < b or c < b: # b is the lowest-index line of the cycle | |
| continue | |
| for d in adj[c].get(b, ()): | |
| if d <= b or d == a or d == c: | |
| continue | |
| if (min(c, a), max(c, a)) in seg[d] and (min(d, b), max(d, b)) in seg[a]: | |
| cyc.append((a, b, c, d)) | |
| if not cyc: | |
| return [], prof | |
| cyc = np.asarray(cyc, np.int64) | |
| a, b, c, d = cyc[:, 0], cyc[:, 1], cyc[:, 2], cyc[:, 3] | |
| mgn = cfg.border_margin_px | |
| img_area = float(h * w) | |
| X = np.stack([PX[a, b], PX[b, c], PX[c, d], PX[d, a]], 1) # (M,4) corner P_ab, P_bc, P_cd, P_da | |
| Y = np.stack([PY[a, b], PY[b, c], PY[c, d], PY[d, a]], 1) | |
| keep = np.isfinite(X).all(1) & np.isfinite(Y).all(1) | |
| keep &= (X >= -mgn).all(1) & (X <= w - 1 + mgn).all(1) & (Y >= -mgn).all(1) & (Y <= h - 1 + mgn).all(1) | |
| ex, ey = np.roll(X, -1, 1) - X, np.roll(Y, -1, 1) - Y | |
| cr = ex * np.roll(ey, -1, 1) - ey * np.roll(ex, -1, 1) | |
| keep &= ((cr > 1e-6).all(1) | (cr < -1e-6).all(1)) | |
| ln = np.hypot(ex, ey) | |
| keep &= (ln >= cfg.min_side_px).all(1) | |
| area = 0.5 * np.abs((X * np.roll(Y, -1, 1) - np.roll(X, -1, 1) * Y).sum(1)) / img_area | |
| keep &= (area >= cfg.min_area_frac) & (area <= cfg.max_area_frac) | |
| asp = np.maximum((ln[:, 0] + ln[:, 2]) / np.maximum(ln[:, 1] + ln[:, 3], 1e-6), | |
| (ln[:, 1] + ln[:, 3]) / np.maximum(ln[:, 0] + ln[:, 2], 1e-6)) | |
| keep &= asp <= cfg.max_aspect | |
| if cfg.shape == "rect": | |
| dot = ex * np.roll(ex, -1, 1) + ey * np.roll(ey, -1, 1) | |
| cosang = -dot / np.maximum(ln * np.roll(ln, -1, 1), 1e-9) # interior angle at corner k+1 | |
| keep &= (cosang <= np.cos(np.deg2rad(cfg.corner_angle_deg[0]))).all(1) | |
| keep &= (cosang >= np.cos(np.deg2rad(cfg.corner_angle_deg[1]))).all(1) | |
| k = np.nonzero(keep)[0] | |
| out = [] | |
| if k.size == 0: | |
| return out, prof | |
| X, Y, a, b, c, d = X[k], Y[k], a[k], b[k], c[k], d[k] | |
| area_px = area[k] * img_area | |
| cb, b1, b2 = _cov_vec(cum, S, b, T[b, a], T[b, c], cfg.trim) | |
| cc_, c1, c2 = _cov_vec(cum, S, c, T[c, b], T[c, d], cfg.trim) | |
| cd, d1, d2 = _cov_vec(cum, S, d, T[d, c], T[d, a], cfg.trim) | |
| ca, a1, a2 = _cov_vec(cum, S, a, T[a, d], T[a, b], cfg.trim) | |
| cov = np.stack([cb, cc_, cd, ca], 1) | |
| good = (cov.min(1) >= cfg.min_side_cov) & (cov.mean(1) >= cfg.min_mean_cov) | |
| gi = np.nonzero(good)[0] | |
| if gi.size > cfg.max_eval: # best by coverage (+ area as tie-breaker) | |
| key = cov.mean(1)[gi] + 1e-7 * area_px[gi] | |
| gi = gi[np.argsort(-key)[: cfg.max_eval]] | |
| for j in gi: | |
| info = [(int(b[j]), int(b1[j]), int(b2[j])), (int(c[j]), int(c1[j]), int(c2[j])), | |
| (int(d[j]), int(d1[j]), int(d2[j])), (int(a[j]), int(a1[j]), int(a2[j]))] | |
| out.append((float(cov[j].mean()), float(cov[j].min()), | |
| np.stack([X[j], Y[j]], 1), cov[j].copy(), info, float(area_px[j]))) | |
| return out, prof | |
| def _shared_sides(parent: np.ndarray, child: np.ndarray, tol_px: float = 3.0, tol_deg: float = 3.0) -> int: | |
| """How many sides of 'child' lie on the same line as a side of 'parent'.""" | |
| def sides(c): | |
| p, q = c, np.roll(c, -1, axis=0) | |
| d = q - p | |
| ln = np.maximum(np.hypot(d[:, 0], d[:, 1]), 1e-9) | |
| n = np.stack([-d[:, 1], d[:, 0]], 1) / ln[:, None] | |
| return n, (p + q) / 2 | |
| np_, mp = sides(parent) | |
| nc, mc = sides(child) | |
| cosmin = np.cos(np.deg2rad(tol_deg)) | |
| n_shared = 0 | |
| for k in range(4): | |
| par = np.abs(nc[k][None] @ np_.T)[0] >= cosmin # parent sides parallel | |
| dist = np.abs((mc[k][None] - mp) @ np_.T).diagonal() if False else np.abs(((mc[k] - mp) * np_).sum(1)) | |
| if (par & (dist <= tol_px)).any(): | |
| n_shared += 1 | |
| return n_shared | |
| def _select(scored, cfg: QuadConfig): | |
| """Final selection in 3 steps. | |
| 1. NMS only on near-duplicates (IoU > 0.7): the nested shapes also survive. | |
| 2. 'Stitched group': a box that contains >=2 disjoint children covering >=group_cover of it | |
| is NOT an object but the space between nearby objects (3 paintings in a row): the | |
| parent is dropped and the children stay. A frame with ONE inner window is not a group. | |
| 3. nested=False: among the survivors only the OUTER contour is kept (>=70% contained in | |
| another one -> dropped); nested=True: all of them are kept. | |
| The real intersections (cv2) are computed only when the bounding box does not rule them out.""" | |
| def bbox(c): | |
| return c[:, 0].min(), c[:, 1].min(), c[:, 0].max(), c[:, 1].max() | |
| def bb_inter(p, q): | |
| w = min(p[2], q[2]) - max(p[0], q[0]); h = min(p[3], q[3]) - max(p[1], q[1]) | |
| return max(w, 0.0) * max(h, 0.0) | |
| U, UB, UA, ALT = [], [], [], [] | |
| for item in scored: | |
| c = item[1]; bb = bbox(c); ar = _poly_area(c) | |
| dup_idx = -1 | |
| for k, (ob, oa, o) in enumerate(zip(UB, UA, U)): | |
| bi = bb_inter(bb, ob) | |
| if bi / max(ar + oa - bi, 1e-6) <= 0.7: # upper bound of the IoU | |
| continue | |
| if _overlap(c, o[1])[0] > 0.7: | |
| dup_idx = k | |
| break | |
| if dup_idx < 0: | |
| U.append(item); UB.append(bb); UA.append(ar); ALT.append([]) | |
| else: | |
| ALT[dup_idx].append((item, bb, ar)) # near-duplicate: alternative for the kept item | |
| if len(U) >= 4 * cfg.max_quads: | |
| break | |
| if cfg.prefer_outer and not cfg.nested: | |
| # among the near-duplicates prefer the OUTER contour (the outer border of a frame instead | |
| # of the inner border), but only if the larger alternative does not invade another selected | |
| # object: otherwise it is a rectangle widening into the space between two paintings | |
| for k in range(len(U)): | |
| best = None | |
| for item, bb, ar in ALT[k]: | |
| if ar <= UA[k] or item[0] < U[k][0] - cfg.prefer_outer_tol: | |
| continue | |
| if _overlap(U[k][1], item[1])[1] < 0.9: | |
| continue | |
| # an invaded object: it overlaps without being (almost) fully inside the alternative | |
| if any(m != k and bb_inter(bb, UB[m]) > 0 and _overlap(item[1], U[m][1])[0] > 0.02 | |
| and _overlap(U[m][1], item[1])[1] < 0.7 for m in range(len(U))): | |
| continue | |
| if best is None or ar > best[2]: | |
| best = (item, bb, ar) | |
| if best is not None: | |
| U[k], UB[k], UA[k] = best | |
| n = len(U) | |
| area = UA | |
| inside = np.zeros((n, n), bool) # inside[i, j] = i is >=70% contained in j | |
| for i in range(n): | |
| for j in range(n): | |
| if i != j and area[i] < area[j] and bb_inter(UB[i], UB[j]) / max(area[i], 1e-6) >= cfg.nms_contain: | |
| inside[i, j] = _overlap(U[i][1], U[j][1])[1] >= cfg.nms_contain | |
| drop = set() | |
| # 'stretched' boxes: they extend a better supported child sharing >=2 collinear sides | |
| # (rectangle crossing the space between two aligned objects, or 'stitching' two objects) | |
| for j in range(n): | |
| for i in range(n): | |
| if inside[i, j] and U[i][0] >= U[j][0] + cfg.stretch_margin \ | |
| and _shared_sides(U[j][1], U[i][1]) >= 2: | |
| drop.add(j) | |
| break | |
| for j in range(n): | |
| if j in drop: | |
| continue | |
| kids = [i for i in range(n) if inside[i, j] and i not in drop] | |
| if len(kids) < 2: | |
| continue | |
| kids = [i for i in kids if not any(inside[i, k] for k in kids if k != i)] # direct kids | |
| disj = [] | |
| for i in sorted(kids, key=lambda t: -area[t]): | |
| if all(bb_inter(UB[i], UB[k]) <= 0 or _overlap(U[i][1], U[k][1])[0] < 0.1 for k in disj): | |
| disj.append(i) | |
| if len(disj) >= 2 and sum(area[i] for i in disj) >= cfg.group_cover * area[j]: | |
| drop.add(j) | |
| keep = [i for i in range(n) if i not in drop] | |
| if not cfg.nested: | |
| keep = [i for i in keep if not any(inside[i, j] for j in keep if j != i)] | |
| return [U[i] for i in keep][: cfg.max_quads] | |
| def extract_quads(edge: np.ndarray, cfg: QuadConfig = None, refine: bool = True, | |
| return_lines: bool = False, gray: np.ndarray = None): | |
| """edge: (H,W) float in [0,1]. gray: (H,W) grayscale image, same frame as the edge map; if | |
| missing there is no contrast test and the stitched rectangles are NOT filtered well. Returns | |
| a list of dict {corners(4,2) clockwise from min(x+y), score, edge_cov(4), contrast_cov(4), | |
| polarity, area, refine_shift_px} sorted by score; optionally also the lines.""" | |
| cfg = cfg or QuadConfig() | |
| edge = np.asarray(edge, dtype=np.float32) | |
| h, w = edge.shape | |
| k = int(2 * round(cfg.band_px) + 1) | |
| near = cv2.dilate((edge > cfg.edge_thresh).astype(np.uint8), | |
| cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))) > 0 | |
| orient = _orientation_fields(edge, cfg) | |
| lines = _refit_lines(edge, hough_lines(edge, cfg, orient), cfg)[: cfg.n_lines] | |
| if len(lines) >= 4: | |
| cands, prof = _enumerate(lines, near, orient, h, w, cfg) | |
| else: | |
| cands, prof = [], None | |
| # pre-sort by coverage (+ area as a tie-breaker): only the best go to the slow checks | |
| cands.sort(key=lambda z: -(z[0] + 1e-7 * z[5])) | |
| cands = cands[: cfg.max_eval] | |
| scored = [] | |
| for mean_cov, _, c, sc, info, _ar in cands: | |
| # largest gap per side: a rectangle "stitched" from nearby objects has a gap = the space | |
| # between the objects, even if the mean coverage stays high | |
| if cfg.max_gap_frac > 0 and any(_longest_gap(prof[L], i1, i2) > cfg.max_gap_frac * max(i2 - i1, 1) | |
| for L, i1, i2 in info): | |
| continue | |
| c = canonical_order(c) | |
| sc = [float(s) for s in sc] | |
| if gray is not None and cfg.min_contrast_cov > 0: | |
| cc, pol = _contrast_support(gray, c, cfg) | |
| if min(cc) < cfg.min_contrast_cov: | |
| continue | |
| rank = float(np.mean(cc)) + 0.5 * mean_cov + 1e-7 * _ar | |
| else: | |
| cc, pol = [], 0 | |
| rank = mean_cov + 1e-7 * _ar | |
| scored.append((rank, c, sc, cc, pol)) | |
| scored.sort(key=lambda z: -z[0]) | |
| chosen = _select(scored, cfg) | |
| res = [] | |
| for rank, c, sc, cc, pol in chosen: | |
| shift = 0.0 | |
| if refine: | |
| c, shift = _refine(edge, c, cfg) | |
| c = canonical_order(c) | |
| res.append({"corners": c, "score": float(rank), "edge_cov": [float(s) for s in sc], | |
| "contrast_cov": [float(s) for s in cc], "polarity": int(pol), | |
| "area": _poly_area(c), "refine_shift_px": shift}) | |
| return (res, lines) if return_lines else res | |
| def hybrid_edge_map(edge_prob: np.ndarray, gray_u8: np.ndarray, canny_lo: int = 20, | |
| canny_hi: int = 53, blur_sigma: float = 1.2, | |
| valid: Optional[np.ndarray] = None) -> np.ndarray: | |
| """Union (max) of the network edge_prob and of a classical Canny on the same image. | |
| Measured reason: on the real domain the network loses the low-contrast sides, and without | |
| a line for that side the candidate is never born. valid: (H,W) bool mask of the area holding | |
| real image (outside = pad border): there the edges are zeroed.""" | |
| canny = cv2.Canny(cv2.GaussianBlur(gray_u8, (0, 0), blur_sigma), canny_lo, canny_hi) | |
| out = np.maximum(np.asarray(edge_prob, dtype=np.float32), canny.astype(np.float32) / 255.0) | |
| if valid is not None: | |
| out = out * valid.astype(np.float32) | |
| return out | |
| def extract_quads_hybrid(edge_prob: np.ndarray, image_bgr: np.ndarray, cfg: QuadConfig = None, | |
| valid: Optional[np.ndarray] = None, **kw): | |
| """Shortcut: hybrid map + contrast test on the image (same frame as edge_prob).""" | |
| gray_u8 = image_bgr if image_bgr.ndim == 2 else cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) | |
| edge = hybrid_edge_map(edge_prob, gray_u8, valid=valid) | |
| return extract_quads(edge, cfg or QuadConfig.hybrid(), gray=gray_u8.astype(np.float32), **kw) | |