hlt-geometric-vision / code /hlt_postproc_quads.py
tailot's picture
Build release: ONNX pipeline + technical documentation + public-domain examples (scripts/build_hf_release.py)
6464353 verified
Raw History Blame Contribute Delete
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
@dataclass
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)
@classmethod
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)