| |
| """Rectify a photo of an ECG sheet or strip: a flat sheet with a square grid at 10 px/mm. |
| |
| python3 rectify.py photo.jpg flat.jpg [--illum] [--illum-paper {white,keep}] [--redshift] |
| |
| Outputs: flat.jpg (the flat sheet), flat.json (status, self-check, scale, timings), |
| flat-check.jpg (the photo with the mask, working area, crop and grid nodes drawn). |
| |
| Statuses: ok, warn, fail, no_grid (no 5 mm grid: masked crop without rectification), |
| too_small (sheet < 50 x 30 mm), not_applicable (the grid falls into several separate pieces). |
| |
| Pipeline: [redshift] -> mask -> crop -> photo scale -> resize to 5.5 px/mm -> grid nodes -> |
| scope check -> offsets network on windows -> fusion into one map -> node indices (GNN) -> |
| global correction (cubic polynomial or homography) -> local correction by nodes -> flat sheet. |
| """ |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| import traceback |
| from contextlib import contextmanager |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
| from ecglib import correction, fusion, illum, lattice, nets, nodemap, redshift, render, scale, windows |
|
|
| THREADS = 2 |
| OUTPUT_PX_PER_MM = 10.0 |
| TARGET = scale.TARGET_PX_PER_MM |
| MASK_THRESHOLD = 0.5 |
| CROP_PAD = 0.03 |
| HULL_MIN_ISLAND = 0.02 |
| EDGE_KERNEL = 15 |
| TOO_SMALL_MM = (50.0, 30.0) |
| MIN_PIECE = 0.05 |
| WINDOW_SCALES = (0.9, 1.0, 1.1) |
| WINDOW_OVERLAP = 0.5 |
| WINDOW_KEEP = 0.5 |
| MESH_STEP = 16.0 |
| SMOOTHNESS = 1.0 |
| MIN_NODES = 40 |
| MAX_SHEET_MM = 450.0 |
| MAX_POINTS = 5000 |
| OK_MM, WARN_MM = 0.5, 1.0 |
| NODE_MESH_STEP = 8.0 |
| NODE_OUTLIER_MAX = 0.03 |
| NODE_OK_MM, NODE_WARN_MM = 0.2, 0.4 |
|
|
|
|
| class Stop(Exception): |
| """Early end of the pipeline with a final status; flat / check: images to output instead (or None).""" |
|
|
| def __init__(self, status, reason, flat=None, check=None): |
| super().__init__(reason) |
| self.status, self.reason, self.flat, self.check = status, reason, flat, check |
|
|
|
|
| @dataclass |
| class Crop: |
| photo: np.ndarray |
| mask: np.ndarray |
| x0: int |
| y0: int |
| x1: int |
| y1: int |
|
|
|
|
| @dataclass |
| class Prepared: |
| image: np.ndarray |
| mask: np.ndarray |
| hull: np.ndarray |
| factor: float |
| px_per_mm: float |
| verified: bool |
|
|
|
|
| @contextmanager |
| def timer(report, name): |
| t = time.perf_counter() |
| try: |
| yield |
| finally: |
| report.setdefault("times", {})[name] = round(time.perf_counter() - t, 2) |
|
|
|
|
| def rounded(d): |
| return {k: round(v, 3) if isinstance(v, float) else v for k, v in d.items()} |
|
|
|
|
| def to_original(crop, prep, xy): |
| return np.stack([crop.x0 + (xy[:, 0] + 0.5) / prep.factor - 0.5, |
| crop.y0 + (xy[:, 1] + 0.5) / prep.factor - 0.5], 1) |
|
|
|
|
| def draw_contours(image, mask, color, width, shift=(0, 0)): |
| contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| cv2.drawContours(image, [c + np.array(shift) for c in contours], -1, color, width) |
|
|
|
|
| def hull_of(mask): |
| """Convex hull of the mask without small islands: blank paper next to the signal is bent too.""" |
| n, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8)) |
| areas = stats[1:, cv2.CC_STAT_AREA] |
| keep = np.isin(labels, 1 + np.nonzero(areas >= HULL_MIN_ISLAND * max(1, areas.sum()))[0]) |
| points = cv2.findNonZero(keep.astype(np.uint8)) |
| out = np.zeros(mask.shape, np.uint8) |
| if points is not None: |
| cv2.fillConvexPoly(out, cv2.convexHull(points), 1) |
| return out > 0 |
|
|
|
|
| def scale_note(status, reason, verified): |
| if not verified: |
| if status == "ok": |
| status = "warn" |
| reason += "; scale not verified (cell = 5 mm not confirmed): mm may be wrong, e.g. x5" |
| return status, reason |
|
|
|
|
| def step_mask(models, photo, report): |
| mask = nets.mask_net(models, photo) > MASK_THRESHOLD |
| report["mask_fraction"] = float(mask.mean()) |
| if not mask.any(): |
| raise Stop("fail", "empty mask") |
| return mask |
|
|
|
|
| def step_crop_to_mask(photo, mask, report): |
| ys, xs = np.nonzero(mask) |
| pad = int(CROP_PAD * max(photo.shape[:2])) |
| x0, y0 = max(0, xs.min() - pad), max(0, ys.min() - pad) |
| x1, y1 = min(photo.shape[1], xs.max() + pad + 1), min(photo.shape[0], ys.max() + pad + 1) |
| report["crop"] = [int(x0), int(y0), int(x1), int(y1)] |
| return Crop(photo[y0:y1, x0:x1], mask[y0:y1, x0:x1], int(x0), int(y0), int(x1), int(y1)) |
|
|
|
|
| def step_find_scale(models, photo, mask, crop, report): |
| ppm, centre, tried, verified = scale.find_scale_in_mask(models, crop.photo, crop.mask) |
| report.update(scale_tries=tried, scale_verified=verified) |
| if ppm is None: |
| flat = render.crop_to(np.where(mask[..., None], photo, 0), mask) |
| raise Stop("no_grid", "no 5 mm lattice in the mask", flat=flat) |
| report.update(photo_px_per_mm=ppm, scale_centre=centre) |
| return ppm, verified |
|
|
|
|
| def step_prepare(crop, ppm, verified, report): |
| factor = TARGET / ppm |
| image = scale.resize(crop.photo, factor) |
| h, w = image.shape[:2] |
| mask = cv2.resize(crop.mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST) > 0 |
| hull = hull_of(mask) |
| report["prepared"] = [w, h] |
| report["hull_extra_fraction"] = float((hull & ~mask).sum() / max(1, mask.sum())) |
| return Prepared(image, mask, hull, factor, ppm, verified) |
|
|
|
|
| def step_check_size(prep, report): |
| ys, xs = np.nonzero(prep.mask) |
| size_mm = sorted([(xs.max() - xs.min() + 1) / TARGET, (ys.max() - ys.min() + 1) / TARGET], reverse=True) |
| report["sheet_size_mm"] = [round(float(v), 1) for v in size_mm] |
| if size_mm[0] > MAX_SHEET_MM: |
| raise Stop("fail", "sheet {:.0f} mm long (> {:.0f} mm): photo scale is wrong".format(size_mm[0], MAX_SHEET_MM)) |
| if prep.verified and size_mm[0] < TOO_SMALL_MM[0] and size_mm[1] < TOO_SMALL_MM[1]: |
| raise Stop("too_small", "sheet {:.0f} x {:.0f} mm < {:.0f} x {:.0f} mm".format( |
| size_mm[0], size_mm[1], *TOO_SMALL_MM)) |
|
|
|
|
| def step_detect_nodes(models, prep, report): |
| """-> (nodes inside the mask, nodes inside the working area); both slightly dilated to keep edge nodes.""" |
| h, w = prep.image.shape[:2] |
| xy = nets.detector_net(models, prep.image)[:, :2] |
| iy = np.clip(np.round(xy[:, 1]).astype(int), 0, h - 1) |
| ix = np.clip(np.round(xy[:, 0]).astype(int), 0, w - 1) |
| kernel = np.ones((EDGE_KERNEL, EDGE_KERNEL), np.uint8) |
| in_mask = (cv2.dilate(prep.mask.astype(np.uint8), kernel) > 0)[iy, ix] |
| in_hull = (cv2.dilate(prep.hull.astype(np.uint8), kernel) > 0)[iy, ix] |
| report["detector_points"] = int(len(xy)) |
| if max(in_mask.sum(), in_hull.sum()) > MAX_POINTS: |
| raise Stop("fail", "{} detector points (> {}): wrong scale?".format(int(max(in_mask.sum(), in_hull.sum())), MAX_POINTS)) |
| return xy[in_mask], xy[in_hull] |
|
|
|
|
| def step_check_scope(photo, mask, crop, prep, xy_mask, report): |
| """The sheet must fit a network window and the grid must be one piece (not strips in frames).""" |
| h, w = prep.image.shape[:2] |
| if min(w, h) < windows.CROP_SIZE: |
| raise Stop("fail", "sheet too small for a 224 px window") |
| grid = lattice.grid_area(xy_mask, (h, w)) |
| report["grid_fraction_of_mask"] = float((prep.mask & grid).sum() / max(1, prep.mask.sum())) |
| _, _, stats, _ = cv2.connectedComponentsWithStats(grid.astype(np.uint8)) |
| areas = stats[1:, cv2.CC_STAT_AREA].astype(float) |
| report["grid_pieces"] = int((areas >= MIN_PIECE * max(1.0, areas.sum())).sum()) |
| if report["grid_pieces"] >= 2: |
| check = photo.copy() |
| line = max(2, photo.shape[1] // 400) |
| draw_contours(check, mask, (255, 0, 0), line) |
| grid_crop = cv2.resize(grid.astype(np.uint8), (crop.photo.shape[1], crop.photo.shape[0]), |
| interpolation=cv2.INTER_NEAREST) |
| draw_contours(check, grid_crop, (0, 0, 255), line, shift=(crop.x0, crop.y0)) |
| raise Stop("not_applicable", "grid in {} separate pieces, covers {:.0%} of the mask: not one sheet grid " |
| "(strips in frames?)".format(report["grid_pieces"], report["grid_fraction_of_mask"]), check=check) |
|
|
|
|
| def step_rectify_windows(models, prep, xy_hull, report): |
| """Offsets network on windows inside the working area and on the grid.""" |
| h, w = prep.image.shape[:2] |
| grid = lattice.grid_area(xy_hull, (h, w)) |
| use = prep.hull & grid if grid.sum() > 0.1 * prep.mask.sum() else prep.hull |
| all_windows = windows.windows_for(w, h, WINDOW_SCALES, WINDOW_OVERLAP) |
| keep = windows.window_inside(all_windows, use) >= WINDOW_KEEP |
| chosen = [win for win, k in zip(all_windows, keep) if k] |
| report.update(windows_all=int(len(all_windows)), windows_kept=len(chosen)) |
| if len(chosen) < 4: |
| raise Stop("fail", "too few windows in the mask") |
| return chosen, nets.rectify_net(models, prep.image, chosen) |
|
|
|
|
| def step_fuse(prep, chosen, offsets, report): |
| h, w = prep.image.shape[:2] |
| mesh, nodes, report["window_agreement_mm"] = fusion.fuse_windows(w, h, chosen, offsets, MESH_STEP, SMOOTHNESS) |
| return mesh, nodes |
|
|
|
|
| def step_index_nodes(models, xy, report): |
| """Lattice indices by the GNN; the largest connected piece is kept.""" |
| pts, ij = np.zeros((0, 2)), np.zeros((0, 2), int) |
| if len(xy) >= MIN_NODES: |
| ij_all, component = nets.lattice_gnn(models, xy) |
| keep = component >= 0 |
| if keep.any(): |
| main = component == np.bincount(component[keep]).argmax() |
| pts, ij = xy[main], ij_all[main] |
| report["detector_nodes"] = int(len(pts)) |
| return pts, ij |
|
|
|
|
| def step_correct(mesh, network_nodes, pts, ij, verified, report): |
| """Global correction: network + cubic polynomial or a homography, whichever has the lower self-check.""" |
| nodes, status, reason = network_nodes, "warn", "too few detector nodes: no correction, no self-check" |
| if len(pts) >= MIN_NODES: |
| check = correction.holdout_network(mesh, network_nodes, pts, ij) |
| check_h = correction.holdout_homography(mesh, pts, ij) |
| report["selfcheck_network_mm"] = rounded(check) |
| if check_h is not None: |
| report["selfcheck_homography_mm"] = rounded(check_h) |
| apply = correction.fit_correction(mesh.sample(network_nodes, pts), ij) |
| h_nodes = correction.homography_nodes(mesh, pts, ij) |
| if h_nodes is not None and check_h is not None and check_h["median_mm"] < check["median_mm"]: |
| nodes, check, report["method"] = h_nodes, check_h, "homography" |
| else: |
| nodes, report["method"] = apply(network_nodes), "network+poly3" |
| report["selfcheck_holdout_mm"] = rounded(check) |
| med = check["median_mm"] |
| status = "ok" if med <= OK_MM else "warn" if med <= WARN_MM else "fail" |
| reason = "hold-out median {:.2f} mm".format(med) |
| status, reason = scale_note(status, reason, verified) |
| return nodes, status, reason |
|
|
|
|
| def map_from_half(mesh, network_nodes, pts, ij, fit, method): |
| """Global map fitted on the nodes `fit` only (for the self-check).""" |
| if method == "homography": |
| h_nodes = correction.homography_nodes(mesh, pts[fit], ij[fit]) |
| if h_nodes is not None: |
| return lambda p: mesh.sample(h_nodes, p) |
| apply = correction.fit_correction(mesh.sample(network_nodes, pts[fit]), ij[fit]) |
| return lambda p: apply(mesh.sample(network_nodes, p)) |
|
|
|
|
| def step_node_correction(prep, mesh, nodes, network_nodes, pts, ij, verified, status, reason, report): |
| """Local correction by the nodes; skipped when more than NODE_OUTLIER_MAX of the nodes are outliers.""" |
| report["node_correction"] = "off" |
| if len(pts) < MIN_NODES: |
| return mesh, nodes, status, reason |
| apply, outliers = nodemap.correct(mesh.sample(nodes, pts), ij) |
| fraction = outliers / len(pts) |
| report["node_outliers"] = int(outliers) |
| report["node_outlier_fraction"] = round(fraction, 4) |
| if fraction > NODE_OUTLIER_MAX: |
| report["node_correction"] = "skipped" |
| if status == "ok": |
| status = "warn" |
| reason += "; {:.1%} of grid nodes are outliers (> {:.0%}): grid found with errors, node correction not applied".format( |
| fraction, NODE_OUTLIER_MAX) |
| return mesh, nodes, status, reason |
| fit, test = correction.halves(ij) |
| base = map_from_half(mesh, network_nodes, pts, ij, fit, report.get("method")) |
| apply_half, _ = nodemap.correct(base(pts[fit]), ij[fit]) |
| check = correction.score(apply_half(base(pts[test])), ij[test]) |
| report["selfcheck_global_mm"] = report.get("selfcheck_holdout_mm") |
| report["selfcheck_holdout_mm"] = rounded(check) |
| report["method"] = str(report.get("method")) + "+nodes" |
| report["node_correction"] = "applied" |
| med = check["median_mm"] |
| status = "ok" if med <= NODE_OK_MM else "warn" if med <= NODE_WARN_MM else "fail" |
| reason = "hold-out median {:.2f} mm (with node correction)".format(med) |
| status, reason = scale_note(status, reason, verified) |
| h, w = prep.image.shape[:2] |
| fine = fusion.Mesh(w, h, NODE_MESH_STEP) |
| fine_px = np.stack(np.meshgrid(np.arange(fine.nx) * fine.step, np.arange(fine.ny) * fine.step), -1).reshape(-1, 2) |
| return fine, apply(mesh.sample(nodes, fine_px)), status, reason |
|
|
|
|
| def step_render(photo, crop, prep, mesh, nodes, report): |
| source = photo[crop.y0:crop.y1, crop.x0:crop.x1] |
| map_x, map_y = render.inverse_map(source.shape, prep.factor, mesh, nodes, OUTPUT_PX_PER_MM) |
| flat = cv2.remap(source, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0)) |
| flat_mask = cv2.remap(crop.mask.astype(np.uint8) * 255, map_x, map_y, cv2.INTER_NEAREST, |
| borderMode=cv2.BORDER_CONSTANT, borderValue=0) > 0 |
| report["output_px_per_mm"] = OUTPUT_PX_PER_MM |
| return render.crop_to(flat, flat_mask) |
|
|
|
|
| def draw_check(photo, mask, crop, prep, pts): |
| """Mask (blue), working area (yellow), crop (green), grid nodes (red crosses).""" |
| check = photo.copy() |
| w = photo.shape[1] |
| draw_contours(check, mask, (255, 0, 0), max(2, w // 400)) |
| cv2.rectangle(check, (crop.x0, crop.y0), (crop.x1 - 1, crop.y1 - 1), (0, 200, 0), max(1, w // 600)) |
| hull = cv2.resize(prep.hull.astype(np.uint8), (crop.photo.shape[1], crop.photo.shape[0]), |
| interpolation=cv2.INTER_NEAREST) |
| draw_contours(check, hull, (0, 220, 255), max(2, w // 500), shift=(crop.x0, crop.y0)) |
| for x, y in to_original(crop, prep, pts): |
| cv2.drawMarker(check, (int(x), int(y)), (0, 0, 255), cv2.MARKER_CROSS, max(6, w // 150), max(1, w // 800)) |
| return check |
|
|
|
|
| def rectify_photo(models, photo, use_redshift=False): |
| """BGR photo -> (flat sheet or None, check image or None, report). |
| With use_redshift the networks see the recoloured photo; the outputs are built from the original.""" |
| report = {"redshift": False} |
| t0 = time.perf_counter() |
| flat = check = None |
| try: |
| seen = photo |
| if use_redshift: |
| with timer(report, "redshift"): |
| report["redshift"] = True |
| seen = redshift.redshift(photo) |
| with timer(report, "mask"): |
| mask = step_mask(models, seen, report) |
| crop = step_crop_to_mask(seen, mask, report) |
| with timer(report, "scale"): |
| ppm, verified = step_find_scale(models, photo, mask, crop, report) |
| prep = step_prepare(crop, ppm, verified, report) |
| step_check_size(prep, report) |
| with timer(report, "detector"): |
| xy_mask, xy_hull = step_detect_nodes(models, prep, report) |
| step_check_scope(photo, mask, crop, prep, xy_mask, report) |
| with timer(report, "windows"): |
| chosen, offsets = step_rectify_windows(models, prep, xy_hull, report) |
| with timer(report, "fusion"): |
| mesh, network_nodes = step_fuse(prep, chosen, offsets, report) |
| with timer(report, "correction"): |
| pts, ij = step_index_nodes(models, xy_hull, report) |
| nodes, status, reason = step_correct(mesh, network_nodes, pts, ij, verified, report) |
| mesh, nodes, status, reason = step_node_correction(prep, mesh, nodes, network_nodes, pts, ij, |
| verified, status, reason, report) |
| with timer(report, "render"): |
| flat = step_render(photo, crop, prep, mesh, nodes, report) |
| check = draw_check(photo, mask, crop, prep, pts) |
| except Stop as stop: |
| status, reason, flat, check = stop.status, stop.reason, stop.flat, stop.check |
| report.update(status=status, reason=reason) |
| report["times"]["total"] = round(time.perf_counter() - t0, 2) |
| return flat, check, report |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Rectify a photo of an ECG sheet.") |
| parser.add_argument("photo", help="input photo") |
| parser.add_argument("output", help="output image (flat sheet)") |
| parser.add_argument("--illum", action="store_true", help="also write <output>-illum.jpg with flattened illumination") |
| parser.add_argument("--illum-paper", choices=["white", "keep"], default=None, |
| help="paper colour for --illum (default white); implies --illum") |
| parser.add_argument("--redshift", action="store_true", help="recolour dark lines red before the networks (black or grey grids)") |
| args = parser.parse_args() |
| if args.illum_paper is not None: |
| args.illum = True |
| in_path, out_path = Path(args.photo), Path(args.output) |
| photo = cv2.imread(str(in_path), cv2.IMREAD_COLOR) |
| if photo is None: |
| sys.exit("cannot read {}".format(in_path)) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| stem = out_path.with_suffix("") |
| models = nets.load_models(threads=THREADS) |
| try: |
| flat, check, report = rectify_photo(models, photo, use_redshift=args.redshift) |
| except Exception as error: |
| traceback.print_exc() |
| flat, check, report = None, None, {"status": "fail", "reason": repr(error)} |
| report["photo"] = str(in_path) |
| if flat is not None: |
| cv2.imwrite(str(out_path), flat, [cv2.IMWRITE_JPEG_QUALITY, 90]) |
| if args.illum: |
| t = time.perf_counter() |
| cv2.imwrite(str(stem) + "-illum.jpg", illum.flatten(flat, OUTPUT_PX_PER_MM, paper=args.illum_paper or "white"), |
| [cv2.IMWRITE_JPEG_QUALITY, 92]) |
| report["times"]["illum"] = round(time.perf_counter() - t, 2) |
| if check is not None: |
| cv2.imwrite(str(stem) + "-check.jpg", check, [cv2.IMWRITE_JPEG_QUALITY, 85]) |
| Path(str(stem) + ".json").write_text(json.dumps(report, indent=1, default=float, ensure_ascii=False)) |
| selfcheck = report.get("selfcheck_holdout_mm", {}).get("median_mm") |
| ppm = report.get("photo_px_per_mm") |
| print("{}: {} ({})".format(in_path.name, report.get("status"), report.get("reason"))) |
| print(" method: {}, photo scale: {} px/mm, self-check median: {} mm, time: {} s".format( |
| report.get("method"), None if ppm is None else round(ppm, 2), selfcheck, report.get("times", {}).get("total"))) |
| print(" flat sheet: {}".format(out_path if flat is not None else "not produced")) |
| sys.exit(0 if flat is not None else 1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
