Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """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 | |
| class Crop: | |
| photo: np.ndarray | |
| mask: np.ndarray | |
| x0: int | |
| y0: int | |
| x1: int | |
| y1: int | |
| class Prepared: | |
| image: np.ndarray # crop at 5.5 px/mm | |
| mask: np.ndarray | |
| hull: np.ndarray # working area: convex hull of the mask | |
| factor: float # px_5.5 = (px_original + 0.5) * factor - 0.5 | |
| px_per_mm: float | |
| verified: bool # scale confirmed by the 5 mm cell check | |
| 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() | |