188 lines
No EOL
7.5 KiB
Python
188 lines
No EOL
7.5 KiB
Python
"""Generate 1x3 QA screenshots (axial/coronal/sagittal) per volume.
|
||
|
||
Each proc volume is first mapped to a canonical body array (S, P, L) using the
|
||
direction matrix in procmeta (columns = image-axis directions, LPS basis), so
|
||
the panels are correct for any source orientation (LPS axial, LIP coronal,
|
||
PIR sagittal, ...). Radiological display conventions:
|
||
axial: anterior on top, patient's left on right
|
||
coronal: superior on top, patient's left on right
|
||
sagittal: superior on top, anterior on left
|
||
Layout: single row, axial | coronal | sagittal; each panel is scaled so its
|
||
maximal physical dimension (mm) is the same across the row (panels are
|
||
centered vertically, widths follow the physical aspect of each cut).
|
||
Usage:
|
||
python scripts/make_screenshots.py --keys k1 k2 ...
|
||
python scripts/make_screenshots.py --sample 20 [--keys ...]
|
||
Output: data/qa/<patient>/<rest>.png (per-patient, mirrors data/lee_nii)
|
||
"""
|
||
import os
|
||
import sys
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
import argparse
|
||
import json
|
||
import random
|
||
import glob
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
from src.common import ROOT, d, read_nii_img
|
||
|
||
B_L = np.array([1.0, 0, 0])
|
||
B_P = np.array([0.0, 1, 0])
|
||
B_S = np.array([0.0, 0, 1])
|
||
|
||
|
||
def vol_to_body(a, M):
|
||
"""a: array (z,y,x) over image axes; M: 3x3 LPS, columns = axis dirs.
|
||
Returns (body, ok, kof): body = (S, P, L)-indexed array, kof maps each
|
||
body axis to its image axis index (0=x, 1=y, 2=z); kof is None if not
|
||
body-aligned."""
|
||
M = np.asarray(M).reshape(3, 3)
|
||
v = a.transpose(2, 1, 0) # (i0, i1, i2) image axes
|
||
kof, sgn = {}, {}
|
||
for name, b in (("L", B_L), ("P", B_P), ("S", B_S)):
|
||
dots = M.T @ b
|
||
k = int(np.argmax(np.abs(dots)))
|
||
# 0.92 ~= 23 deg: tolerates tilted clinical protocols (e.g. oblique
|
||
# sagittal MP-RAGE ~21 deg from P-S)
|
||
if abs(dots[k]) < 0.92:
|
||
return None, False, None
|
||
kof[name], sgn[name] = k, (1 if dots[k] > 0 else -1)
|
||
out = np.transpose(v, [kof["S"], kof["P"], kof["L"]])
|
||
for ax, name in enumerate(("S", "P", "L")):
|
||
if sgn[name] < 0:
|
||
out = np.flip(out, ax)
|
||
return np.ascontiguousarray(out), True, kof
|
||
|
||
|
||
def panels_from_body(b):
|
||
s, p, l = b.shape
|
||
sm, pm, lm = s // 2, p // 2, l // 2
|
||
axial = b[sm] # (P, L): top=anterior, right=left
|
||
coronal = b[:, pm, :][::-1] # (S-rev, L): top=superior, right=left
|
||
sagittal = b[:, :, lm][::-1] # (S-rev, P): top=superior, left=anterior
|
||
return axial, coronal, sagittal, (sm, pm, lm)
|
||
|
||
|
||
def shot_path(out_dir, key):
|
||
"""Per-patient screenshot path: <out_dir>/<patient>/<rest>.png, where the
|
||
patient id is the 2nd underscore token of the key (lee_/m6_/ntuh_ prefixes)."""
|
||
parts = key.split("_")
|
||
if len(parts) < 3:
|
||
return os.path.join(out_dir, key + ".png")
|
||
return os.path.join(out_dir, parts[1], "_".join(parts[2:]) + ".png")
|
||
|
||
|
||
def screenshot_from_volume(arr, direction, key, out_dir, dpi=110, spacing=None,
|
||
series_desc=None, protocol=None):
|
||
"""Render the 1x3 QA screenshot (axial | coronal | sagittal) from an
|
||
image-axis array (z,y,x), a 3x3 LPS direction matrix (columns =
|
||
image-axis directions) and the image-axis spacing in mm, (x, y, z)
|
||
(None = isotropic 1mm). series_desc / protocol (optional) are shown
|
||
on a second title line. Panel cells follow the physical (mm) aspect
|
||
of each cut. Output: <out_dir>/<patient>/<rest>.png (mirrors the
|
||
data/lee_nii layout). Returns the output path, or None if the volume
|
||
is not body-aligned."""
|
||
if spacing is None:
|
||
spacing = (1.0, 1.0, 1.0)
|
||
body, ok, kof = vol_to_body(arr, direction)
|
||
if not ok:
|
||
return None
|
||
sp = {n: float(spacing[kof[n]]) for n in ("S", "P", "L")}
|
||
lo, hi = np.percentile(body[body > 0], [1, 99.5])
|
||
axial, coronal, sagittal, _ = panels_from_body(body)
|
||
s, p, l = body.shape
|
||
# (image, title, aspect = mm per row / mm per col, n_rows, n_cols)
|
||
panels = [
|
||
(axial, "axial", sp["P"] / sp["L"], p, l),
|
||
(coronal, "coronal", sp["S"] / sp["L"], s, l),
|
||
(sagittal, "sagittal", sp["S"] / sp["P"], s, p),
|
||
]
|
||
max_in, gap_in, m_in, top_in, bot_in = 4.4, 0.25, 0.3, 1.5, 0.2
|
||
# each panel is scaled so its maximal physical dimension spans max_in
|
||
w_in, h_in = [], []
|
||
for _, _, a, nrow, ncol in panels:
|
||
pw, ph = ncol, nrow * a # physical extent along display x / y
|
||
m = max(pw, ph)
|
||
w_in.append(max_in * pw / m)
|
||
h_in.append(max_in * ph / m)
|
||
plot_h = max_in
|
||
fig_w = sum(w_in) + 2 * gap_in + 2 * m_in
|
||
fig_h = plot_h + top_in + bot_in
|
||
fig = plt.figure(figsize=(fig_w, fig_h))
|
||
left = m_in / fig_w
|
||
for (im, title, a, _, _), w, h in zip(panels, w_in, h_in):
|
||
axi = fig.add_axes([left, (bot_in + (plot_h - h) / 2) / fig_h,
|
||
w / fig_w, h / fig_h])
|
||
axi.imshow(im, cmap="gray", vmin=lo, vmax=hi, origin="upper", aspect=a)
|
||
axi.set_title(title, color="w", fontsize=11)
|
||
axi.set_xticks([]); axi.set_yticks([])
|
||
for spine in axi.spines.values():
|
||
spine.set_edgecolor("0.35")
|
||
left += (w + gap_in) / fig_w
|
||
mx = arr.shape[::-1] # (x, y, z)
|
||
tkey = key.split("_", 1)[1] if "_" in key else key # drop dataset prefix (lee_/m6_/ntuh_)
|
||
title = (f"{tkey} {mx[0]}×{mx[1]}×{mx[2]} @ "
|
||
f"{spacing[0]:.3f}×{spacing[1]:.3f}×{spacing[2]:.3f}mm")
|
||
def _clean(s):
|
||
# drop chars the Agg font cannot render (e.g. CJK in old study descs)
|
||
return " ".join(s.split()).encode("ascii", "ignore").decode()
|
||
|
||
sd = _clean(series_desc) if series_desc else ""
|
||
pt = _clean(protocol) if protocol else ""
|
||
info = []
|
||
if sd:
|
||
info.append("series: " + sd[:60])
|
||
if pt and pt[:10] != sd[:10]:
|
||
info.append("protocol: " + pt[:60])
|
||
if info:
|
||
title += "\n" + " ".join(info)
|
||
fig.suptitle(title, color="w", fontsize=12)
|
||
fig.patch.set_facecolor("k")
|
||
out = shot_path(out_dir, key)
|
||
os.makedirs(os.path.dirname(out), exist_ok=True)
|
||
fig.savefig(out, facecolor="k", dpi=dpi)
|
||
plt.close(fig)
|
||
return out
|
||
|
||
|
||
def screenshot(key, out_dir, dpi=110):
|
||
p = os.path.join(d("data/proc"), key + ".nii.gz")
|
||
if not os.path.exists(p):
|
||
return None
|
||
arr = np.asarray(sitk.GetArrayFromImage(read_nii_img(p)), dtype=np.float32)
|
||
meta = json.load(open(os.path.join(d("data/procmeta"), key + ".json")))
|
||
return screenshot_from_volume(arr, meta["direction"], key, out_dir, dpi)
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--keys", nargs="*", default=[])
|
||
ap.add_argument("--sample", type=int, default=0)
|
||
ap.add_argument("--seed", type=int, default=0)
|
||
ap.add_argument("--out", default=d("data/qa"))
|
||
args = ap.parse_args()
|
||
os.makedirs(args.out, exist_ok=True)
|
||
allkeys = [os.path.basename(f)[:-7] for f in glob.glob(os.path.join(d("data/proc"), "*.nii.gz"))
|
||
if not f.endswith("_label.nii.gz")]
|
||
keys = list(args.keys)
|
||
if args.sample:
|
||
rnd = random.Random(args.seed)
|
||
extra = [k for k in allkeys if k not in keys]
|
||
keys += rnd.sample(extra, min(args.sample - len(keys), len(extra)))
|
||
n = 0
|
||
for k in keys:
|
||
r = screenshot(k, args.out)
|
||
if r:
|
||
n += 1
|
||
print("wrote", r, flush=True)
|
||
else:
|
||
print("skip", k, flush=True)
|
||
print(f"done: {n}/{len(keys)} screenshots in {args.out}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |