Ondrej Ciganik <oc@ioproperty.cz> 2026-09-01T19:05:09+02:00
Commit #122 · 1 snippets
scripts/extract_chevrons.py | 242 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 242 insertions(+)
| 1 | """ | ||
| 2 | extract_chevrons.py — extract chevron / gore hatch markings from annotated polygons. | ||
| 3 | |||
| 4 | The dataset-sorter lets a human draw a bounding POLYGON ('markings' in the | ||
| 5 | *_vectors.adjusted.json) around each chevron / gore-hatch zone (where two roads meet). | ||
| 6 | For every kept polygon this extracts the INDIVIDUAL diagonal hatch stripes inside it. | ||
| 7 | |||
| 8 | Model: chevron hatching is a REPEATING set of parallel stripes at one dominant angle | ||
| 9 | (plain hatch) or two mirror angles (chevron arrows). So: | ||
| 10 | 1. mask the bitmap to the polygon, threshold the brightest pixels (paint), | ||
| 11 | 2. connected components -> candidate strokes, | ||
| 12 | 3. split each stroke by per-pixel orientation (structure tensor) so a 'V' becomes its | ||
| 13 | two straight arms, | ||
| 14 | 4. keep only strokes whose angle matches one of the 1-2 dominant hatch angles | ||
| 15 | (this drops clutter; the hatch is prioritised over the bounding side lines), | ||
| 16 | 5. drop long strokes running along the gore axis (the bounding side lines), | ||
| 17 | 6. report the dominant angle(s) and the stripe pitch (spacing) per zone. | ||
| 18 | |||
| 19 | Outputs per tile under <out>/<Abschnitt>/: | ||
| 20 | <tile>_chevrons.json { tile, image_size, chevrons:[{marking_id, polygon_px, | ||
| 21 | polygon_axis_deg, dominant_angles_deg, pitch_px, n_stripes, | ||
| 22 | stripes:[{seg_px:[[x0,y0],[x1,y1]], length_px, width_px, angle_deg}]}] } | ||
| 23 | <tile>_mkN.png QA overlay crop (yellow polygon; stripes coloured by angle group) | ||
| 24 | plus index.html and summary.json. | ||
| 25 | |||
| 26 | Read-only on the dataset. Usage (from repo root): | ||
| 27 | python scripts/extract_chevrons.py --data "<...260624_data - correction_2\\data>" --out out/chevrons | ||
| 28 | """ | ||
| 29 | import argparse, glob, html, json, os, re | ||
| 30 | from pathlib import Path | ||
| 31 | import numpy as np | ||
| 32 | import cv2 | ||
| 33 | from scipy.ndimage import gaussian_filter1d | ||
| 34 | from scipy.signal import find_peaks | ||
| 35 | |||
| 36 | |||
| 37 | def resolve_bmp(ext_dir, tile): | ||
| 38 | for name in (tile, os.path.splitext(tile)[0].replace("_intensity", "") + ".png"): | ||
| 39 | p = os.path.join(ext_dir, name) | ||
| 40 | if os.path.exists(p): | ||
| 41 | return p | ||
| 42 | return None | ||
| 43 | |||
| 44 | |||
| 45 | def orient_map(g, sigma=2.5): | ||
| 46 | """Per-pixel line direction (deg, 0-180) via the structure tensor of the intensity.""" | ||
| 47 | gf = g.astype(np.float32) | ||
| 48 | gx = cv2.Sobel(gf, cv2.CV_32F, 1, 0, ksize=3) | ||
| 49 | gy = cv2.Sobel(gf, cv2.CV_32F, 0, 1, ksize=3) | ||
| 50 | Jxx = cv2.GaussianBlur(gx * gx, (0, 0), sigma) | ||
| 51 | Jyy = cv2.GaussianBlur(gy * gy, (0, 0), sigma) | ||
| 52 | Jxy = cv2.GaussianBlur(gx * gy, (0, 0), sigma) | ||
| 53 | return (np.degrees(0.5 * np.arctan2(2 * Jxy, Jxx - Jyy)) + 90.0) % 180.0 | ||
| 54 | |||
| 55 | |||
| 56 | def _cdiff(a, b): | ||
| 57 | """signed circular difference a-b folded to (-90, 90].""" | ||
| 58 | return (a - b + 90.0) % 180.0 - 90.0 | ||
| 59 | |||
| 60 | |||
| 61 | def _fit(P): | ||
| 62 | c = P.mean(0) | ||
| 63 | _, _, Vt = np.linalg.svd(P - c, full_matrices=False) | ||
| 64 | proj = (P - c) @ Vt[0] | ||
| 65 | perp = (P - c) @ Vt[1] | ||
| 66 | length = float(proj.max() - proj.min()) | ||
| 67 | width = float(perp.max() - perp.min()) + 1.0 | ||
| 68 | p0 = (c + Vt[0] * proj.min()); p1 = (c + Vt[0] * proj.max()) | ||
| 69 | ang = float(np.degrees(np.arctan2(Vt[0][1], Vt[0][0])) % 180) | ||
| 70 | return p0, p1, length, width, ang | ||
| 71 | |||
| 72 | |||
| 73 | def _dominant_angles(strokes, sep=22, maxk=2): | ||
| 74 | if not strokes: | ||
| 75 | return [] | ||
| 76 | h = np.zeros(60) | ||
| 77 | for p0, p1, length, w, ang in strokes: | ||
| 78 | h[int(ang // 3) % 60] += 1 # COUNT-weighted: the hatch is many repeating strokes | ||
| 79 | h = gaussian_filter1d(h, 1.2, mode="wrap") | ||
| 80 | hh = np.concatenate([h, h, h]) | ||
| 81 | pk, _ = find_peaks(hh, distance=5, prominence=max(h.max() * 0.12, 1e-6)) | ||
| 82 | cand = sorted([(hh[p], (p % 60) * 3 + 1.5) for p in pk if 60 <= p < 120], reverse=True) | ||
| 83 | out, hts = [], [] | ||
| 84 | for ht, deg in cand: | ||
| 85 | if all(min(abs(deg - o), 180 - abs(deg - o)) >= sep for o in out): | ||
| 86 | out.append(deg); hts.append(ht) | ||
| 87 | if len(out) >= maxk: | ||
| 88 | break | ||
| 89 | # a real second hatch direction must have substantial support (>=40% of the main one) | ||
| 90 | if len(out) == 2 and hts[1] < 0.4 * hts[0]: | ||
| 91 | out = out[:1] | ||
| 92 | return out or [float(np.median([s[4] for s in strokes]))] | ||
| 93 | |||
| 94 | |||
| 95 | def extract_zone(g, ld, poly, pct=78, min_len=12, min_area=18, split_spread=12, | ||
| 96 | angle_tol=18, lane_tol=18): | ||
| 97 | mask = np.zeros(g.shape, np.uint8) | ||
| 98 | cv2.fillPoly(mask, [poly.astype(np.int32)], 255) | ||
| 99 | inside = g[mask > 0] | ||
| 100 | if inside.size < 80: | ||
| 101 | return [], [], None, 0.0 | ||
| 102 | Pc = poly - poly.mean(0) | ||
| 103 | _, _, pVt = np.linalg.svd(Pc, full_matrices=False) | ||
| 104 | poly_axis = float(np.degrees(np.arctan2(pVt[0][1], pVt[0][0])) % 180) | ||
| 105 | poly_long = float(np.ptp(Pc @ pVt[0])) | ||
| 106 | |||
| 107 | thr = np.percentile(inside, pct) | ||
| 108 | bw = ((g.astype(np.int32) > thr) & (mask > 0)).astype(np.uint8) | ||
| 109 | bw = cv2.morphologyEx(bw, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8)) | ||
| 110 | n, lab, st, _ = cv2.connectedComponentsWithStats(bw, 8) | ||
| 111 | |||
| 112 | raw = [] # (p0,p1,length,width,angle) | ||
| 113 | for i in range(1, n): | ||
| 114 | if st[i, cv2.CC_STAT_AREA] < min_area: | ||
| 115 | continue | ||
| 116 | ys, xs = np.where(lab == i) | ||
| 117 | P = np.column_stack([xs, ys]).astype(float) | ||
| 118 | c = P.mean(0) | ||
| 119 | _, _, Vt = np.linalg.svd(P - c, full_matrices=False) | ||
| 120 | comp_axis = np.degrees(np.arctan2(Vt[0][1], Vt[0][0])) % 180 | ||
| 121 | dirs = ld[ys, xs] | ||
| 122 | delta = _cdiff(dirs, comp_axis) | ||
| 123 | groups = [P] | ||
| 124 | if np.percentile(np.abs(delta), 75) > split_spread and len(P) >= 40: | ||
| 125 | a = P[delta > 4]; b = P[delta < -4] # split a bent V into its two arms | ||
| 126 | if len(a) >= 16 and len(b) >= 16: | ||
| 127 | groups = [a, b] | ||
| 128 | for Q in groups: | ||
| 129 | p0, p1, length, width, ang = _fit(Q) | ||
| 130 | if length < min_len or length / width < 1.6: | ||
| 131 | continue | ||
| 132 | raw.append((p0, p1, length, width, ang)) | ||
| 133 | |||
| 134 | # drop the bounding side lines (long, along the gore axis) BEFORE finding the pattern angle | ||
| 135 | kept = [(p0, p1, length, width, ang) for (p0, p1, length, width, ang) in raw | ||
| 136 | if not (min(abs(ang - poly_axis), 180 - abs(ang - poly_axis)) < lane_tol | ||
| 137 | and length > 0.6 * poly_long)] | ||
| 138 | angles = _dominant_angles(kept) | ||
| 139 | stripes = [s for s in kept | ||
| 140 | if any(min(abs(s[4] - a), 180 - abs(s[4] - a)) <= angle_tol for a in angles)] | ||
| 141 | |||
| 142 | # pitch = median perpendicular spacing within the most-populated angle | ||
| 143 | pitch = None | ||
| 144 | if stripes: | ||
| 145 | a0 = max(angles, key=lambda a: sum(1 for s in stripes | ||
| 146 | if min(abs(s[4] - a), 180 - abs(s[4] - a)) <= angle_tol)) | ||
| 147 | v = np.array([-np.sin(np.radians(a0)), np.cos(np.radians(a0))]) | ||
| 148 | offs = sorted(((np.array(s[0]) + np.array(s[1])) / 2) @ v for s in stripes | ||
| 149 | if min(abs(s[4] - a0), 180 - abs(s[4] - a0)) <= angle_tol) | ||
| 150 | if len(offs) >= 2: | ||
| 151 | d = np.diff(offs); d = d[d > 3] | ||
| 152 | if len(d): | ||
| 153 | pitch = round(float(np.median(d)), 1) | ||
| 154 | return stripes, angles, poly_axis, pitch | ||
| 155 | |||
| 156 | |||
| 157 | def main(): | ||
| 158 | ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | ||
| 159 | ap.add_argument("--data", required=True, help="dataset 'data' dir (has 00_external + 01_interim)") | ||
| 160 | ap.add_argument("--out", default="out/chevrons") | ||
| 161 | ap.add_argument("--pct", type=float, default=78) | ||
| 162 | args = ap.parse_args() | ||
| 163 | |||
| 164 | ext_root = os.path.join(args.data, "00_external", "260611_topdown_tiles") | ||
| 165 | int_root = os.path.join(args.data, "01_interim", "260611_topdown_tiles_markings") | ||
| 166 | out = Path(args.out); out.mkdir(parents=True, exist_ok=True) | ||
| 167 | COLS = [(0, 0, 255), (255, 90, 0), (0, 200, 0)] | ||
| 168 | |||
| 169 | rows, tot_poly, tot_stripe = [], 0, 0 | ||
| 170 | for absn in sorted(os.listdir(int_root)): | ||
| 171 | tdir = os.path.join(int_root, absn, "tiles") | ||
| 172 | if not os.path.isdir(tdir): | ||
| 173 | continue | ||
| 174 | odir = out / absn.replace("260611_", ""); odir.mkdir(parents=True, exist_ok=True) | ||
| 175 | for f in sorted(glob.glob(os.path.join(tdir, "*.adjusted.json"))): | ||
| 176 | o = json.load(open(f, encoding="utf-8")) | ||
| 177 | polys = [m for m in o.get("markings", []) if not m.get("removed")] | ||
| 178 | if not polys: | ||
| 179 | continue | ||
| 180 | bmp = resolve_bmp(os.path.join(ext_root, absn), o["tile"]) | ||
| 181 | if not bmp: | ||
| 182 | print(f" [skip] no bitmap for {o['tile']}"); continue | ||
| 183 | g = cv2.imread(bmp, cv2.IMREAD_GRAYSCALE) | ||
| 184 | ld = orient_map(g) | ||
| 185 | base = cv2.cvtColor(g, cv2.COLOR_GRAY2BGR) | ||
| 186 | stem = os.path.splitext(o["tile"])[0] | ||
| 187 | chevrons = [] | ||
| 188 | for i, m in enumerate(polys): | ||
| 189 | poly = np.array(m["points"], float) | ||
| 190 | stripes, angles, axis, pitch = extract_zone(g, ld, poly, pct=args.pct) | ||
| 191 | tot_poly += 1; tot_stripe += len(stripes) | ||
| 192 | chevrons.append(dict( | ||
| 193 | marking_id=m.get("marking_id"), | ||
| 194 | polygon_px=[[round(x, 1), round(y, 1)] for x, y in poly.tolist()], | ||
| 195 | polygon_axis_deg=round(axis, 1) if axis is not None else None, | ||
| 196 | dominant_angles_deg=[round(a, 1) for a in angles], pitch_px=pitch, | ||
| 197 | n_stripes=len(stripes), | ||
| 198 | stripes=[dict(seg_px=[[round(p0[0], 1), round(p0[1], 1)], | ||
| 199 | [round(p1[0], 1), round(p1[1], 1)]], | ||
| 200 | length_px=round(length, 1), width_px=round(width, 1), | ||
| 201 | angle_deg=round(ang, 1)) | ||
| 202 | for p0, p1, length, width, ang in stripes])) | ||
| 203 | # overlay crop | ||
| 204 | x0, y0 = poly[:, 0].min(), poly[:, 1].min(); x1, y1 = poly[:, 0].max(), poly[:, 1].max() | ||
| 205 | mr = 40 | ||
| 206 | xa, ya = int(max(0, x0 - mr)), int(max(0, y0 - mr)) | ||
| 207 | xb, yb = int(min(g.shape[1], x1 + mr)), int(min(g.shape[0], y1 + mr)) | ||
| 208 | crop = base[ya:yb, xa:xb].copy() | ||
| 209 | cv2.polylines(crop, [poly.astype(np.int32) - [xa, ya]], True, (0, 255, 255), 1, cv2.LINE_AA) | ||
| 210 | for p0, p1, length, width, ang in stripes: | ||
| 211 | gi = min(range(len(angles)), key=lambda k: min(abs(ang - angles[k]), 180 - abs(ang - angles[k]))) if angles else 0 | ||
| 212 | a = (np.array(p0) - [xa, ya]).astype(int); b = (np.array(p1) - [xa, ya]).astype(int) | ||
| 213 | cv2.line(crop, tuple(a), tuple(b), COLS[gi % 3], 2, cv2.LINE_AA) | ||
| 214 | img = f"{stem}_mk{i}.png" | ||
| 215 | cv2.imwrite(str(odir / img), crop) | ||
| 216 | rows.append(dict(absn=absn.replace("260611_", ""), tile=stem, mk=i, | ||
| 217 | img=f"{absn.replace('260611_','')}/{img}", n=len(stripes), | ||
| 218 | ang=[round(a) for a in angles], pitch=pitch)) | ||
| 219 | (odir / f"{stem}_chevrons.json").write_text(json.dumps( | ||
| 220 | dict(tile=o["tile"], abschnitt=absn, image_size=[g.shape[1], g.shape[0]], | ||
| 221 | chevrons=chevrons), indent=1)) | ||
| 222 | |||
| 223 | cards = [] | ||
| 224 | for r in rows: | ||
| 225 | cards.append(f'<div class=card><a href="{html.escape(r["img"])}" target=_blank>' | ||
| 226 | f'<img src="{html.escape(r["img"])}" loading=lazy></a>' | ||
| 227 | f'<div class=m>{r["absn"]}/{r["tile"]} mk{r["mk"]} · ' | ||
| 228 | f'<b>{r["n"]}</b> stripes · {r["ang"]}° · pitch {r["pitch"]}</div></div>') | ||
| 229 | (out / "index.html").write_text( | ||
| 230 | "<!doctype html><meta charset=utf-8><title>chevrons</title><style>" | ||
| 231 | "body{background:#111;color:#ddd;font:13px system-ui;margin:12px}.grid{display:flex;flex-wrap:wrap;gap:10px}" | ||
| 232 | ".card{border:1px solid #444;border-radius:6px;padding:4px;background:#1a1a1a;max-width:360px}" | ||
| 233 | ".card img{max-width:360px;height:auto;display:block}.m{padding:4px;font-size:12px}</style>" | ||
| 234 | f"<h2>Chevron extraction — {tot_poly} polygons, {tot_stripe} stripes</h2>" | ||
| 235 | f"<div class=grid>{''.join(cards)}</div>", encoding="utf-8") | ||
| 236 | (out / "summary.json").write_text(json.dumps(dict(polygons=tot_poly, stripes=tot_stripe, tiles=rows), indent=1)) | ||
| 237 | print(f"{tot_poly} polygons -> {tot_stripe} stripes across {len(set((r['absn'],r['tile']) for r in rows))} tiles") | ||
| 238 | print(f"wrote {out}/index.html + per-tile *_chevrons.json") | ||
| 239 | |||
| 240 | |||
| 241 | if __name__ == "__main__": | ||
| 242 | main() | ||
| 0 |