Xiao Furen
d167c1f7c7
Introduce a robust coordinate transformation system to manage the relationship between original CT space and rotated/standardized segmentation spaces. This includes a new directory hierarchy to separate unrotated crops from rotated outputs and utility functions for geometric mapping. Key changes: - Implement `imaging/transforms.py` to handle bounding box metadata, affine standardization, and coordinate mapping between spaces. - Restructure dataset output: unrotated segmentation files (binary, SDF, ROI, etc.) are now stored in a `<vol>/crop/` subdirectory to distinguish them from `<vol>/rotated/` aligned versions. - Add `level_file_path` utility to abstract file discovery across legacy (top-level) and new (crop-based) directory structures. - Enhance `seg_bone` to capture and export bounding box metadata (`bbox2`, `nn_bbox`, `bbox_orig`) into `transform.json`. - Implement `xfr_cbt_native.py` for mapping screw positions back to original CT space. - Update preprocessing and visualization scripts to support the new directory layout and transformation metadata. - Improve TinyDB metadata migration logic to prevent accidental corruption of existing database structures.
628 lines
No EOL
28 KiB
Python
628 lines
No EOL
28 KiB
Python
#!/home/xfr/.conda/envs/cbt/bin/python
|
||
|
||
"""把每個 lumbar level 的 rotated/{L}_label.nii.gz（旋轉座標，
|
||
1=VBODY 椎體、2=SP 棘突、3=other bone）反旋轉回原始座標，
|
||
把所有 level 合併成「每 patient 一個」.nii.gz，存到 CTSpine1K
|
||
原始 dataset 下面（座標系與 data/ 裡的原始影像檔一致）。
|
||
|
||
做法：
|
||
1) 依 level 重算對齊旋轉 (R, c_xyz)（與 _write_rotated_level 完全
|
||
同路徑：_smd_resampled 有效時以 <0.5 取骨頭 mask，否則 _binary_nn
|
||
以 >0），再從旋轉 label 檔與未旋轉模板檔的 origin 差回復裁切
|
||
起點 fstart（(O_rot - O_tpl) 換元、必須為整數）。
|
||
2) NN 反向映射把 label 映回未旋轉 0.5mm 模板 grid。旋轉是 index 空間
|
||
剛性變換（物理座標系統只平移 fstart），所以換元後只是單純的
|
||
R 旋轉，不需要重新估算任何平面。
|
||
3) std（0.5mm 標準化空間）-> 原始 CT 物理座標的對角仿射 mapping：
|
||
standardize_affine 每翻一個軸就把整份資料平移到世界原點的鏡射面
|
||
（b_i = -(2O_i + d_i·0.5·(L_i-1))，翻軸條件 x: d_x=+1、y: d_y=+1、
|
||
z: d_z=-1），ap_flip（metadata db）再令 y 軸鏡射。全部參數都可從
|
||
原始 CT 幾何 + ap_flip 算出，見 _std_to_orig_affine。
|
||
4) NN 重取樣到原始 CT grid（/mnt/1220/Public/dataset/Spine/CTSpine1K/
|
||
data/<子目錄>/<name>.nii.gz 的幾何），依 level 重編號後合併：
|
||
Lx 的 VBODY = 2x-1、SP = 2x（L1..L6 -> 1..12），0 = 背景；
|
||
rotated label 的 3（other bone）不納入。
|
||
每 level 驗證重算的 (R, c)：主要用正向 identity —— 以 _write_rotated_level
|
||
完全同參數重旋轉未旋轉 check 檔（_smd_resampled 三線性、缺則
|
||
_binary_nn NN），與存檔旋轉檔逐體素比對，mismatch > 1e-4 跳過
|
||
（精確驗證 (R,c) 與該輸出生成版本一致，不受薄結構影響）；無可比檔
|
||
時 fallback 用 round-trip 包含率（旋轉 check 檔反向旋轉後與未旋轉
|
||
bone mask 互相包含於 2-voxel 膨脹內，min 方向 < --min-contain
|
||
跳過該 level）。
|
||
寫檔前對位檢查：合併 label 落在原始 CT 體內（HU > -100）比例 < 0.98
|
||
或 median HU < 100 時判定錯位、跳過。
|
||
|
||
輸出：<dest>/<子目錄>/<name>.nii.gz（uint8），預設 dest =
|
||
/mnt/1220/Public/dataset/Spine/CTSpine1K/label_vbody_sp/，
|
||
檔名與原始影像檔相同。已存在的輸出預設跳過（--force 覆蓋）。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from scipy.ndimage import map_coordinates
|
||
|
||
from imaging.orientation import (best_symmetry_plane, best_upper_endplate_plane)
|
||
from imaging.transforms import level_file_path
|
||
from visualization.res_bone_figure import (compute_normalizing_rotation,
|
||
rotated_grid, rotate_volume_to)
|
||
|
||
DATA_ROOT = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
|
||
DATASET_ROOT = '/mnt/1220/Public/dataset/Spine/CTSpine1K/'
|
||
DEFAULT_DEST = os.path.join(DATASET_ROOT, 'label_vbody_sp')
|
||
# 預設依序掃的世代資料夾（先新後舊；同名 volume 以先處理者為準）
|
||
DEFAULT_ROOTS = [
|
||
'/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3',
|
||
'/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2',
|
||
]
|
||
LUMBAR_LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6')
|
||
|
||
LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
|
||
METADATA_DB = os.path.join(os.path.dirname(os.path.abspath(__file__)),
|
||
'xfr_image_metadata.json')
|
||
|
||
logger = logging.getLogger('xfr_orig_labels')
|
||
|
||
|
||
def load_ap_flip():
|
||
"""metadata db（TinyDB 純 JSON：{'images': {doc_id: doc}}）的
|
||
name -> ap_flip（bool）；db 缺失 / 解析失敗時回 {}。"""
|
||
try:
|
||
with open(METADATA_DB) as f:
|
||
data = json.load(f)
|
||
except (json.JSONDecodeError, OSError) as e:
|
||
logger.warning(f'metadata db {METADATA_DB} 讀取失敗：{e}')
|
||
return {}
|
||
tab = data.get('images') if isinstance(data, dict) else None
|
||
if not isinstance(tab, dict):
|
||
return {}
|
||
return {v['name']: bool(v.get('ap_flip', False))
|
||
for v in tab.values() if isinstance(v, dict) and 'name' in v}
|
||
|
||
|
||
# ---------------------------------------------------------------- log 設定
|
||
def setup_tee(log_path):
|
||
"""console 與 log 檔同時輸出（append、line-buffered）。"""
|
||
class _Tee:
|
||
def __init__(self, console, fh):
|
||
self.console, self.fh, self.buf = console, fh, ''
|
||
|
||
def write(self, data):
|
||
if not data:
|
||
return
|
||
self.buf += data
|
||
while True:
|
||
idx_n, idx_r = self.buf.find('\n'), self.buf.find('\r')
|
||
idx = min([i for i in (idx_n, idx_r) if i != -1], default=-1)
|
||
if idx == -1:
|
||
break
|
||
line, self.buf = self.buf[:idx], self.buf[idx + 1:]
|
||
self.console.write(line + '\n')
|
||
self.fh.write(line + '\n')
|
||
|
||
def flush(self):
|
||
if self.buf:
|
||
line, self.buf = self.buf, ''
|
||
self.console.write(line + '\n')
|
||
self.fh.write(line + '\n')
|
||
self.console.flush()
|
||
self.fh.flush()
|
||
|
||
fh = open(log_path, 'a', buffering=1)
|
||
sys.stdout = _Tee(sys.stdout, fh)
|
||
sys.stderr = _Tee(sys.stderr, fh)
|
||
|
||
|
||
# ---------------------------------------------------------------- 幾何工具
|
||
def _axis_vectors(img):
|
||
"""direction 為純對角（±1）時回傳 (signs (x,y,z), spacing)，否則 None。"""
|
||
d = np.asarray(img.GetDirection(), dtype=float).reshape(3, 3)
|
||
diag = np.diag(d)
|
||
if not (np.allclose(d, np.diag(diag), atol=1e-6)
|
||
and np.allclose(np.abs(diag), 1.0, atol=1e-6)):
|
||
return None
|
||
return diag, np.asarray(img.GetSpacing(), dtype=float)
|
||
|
||
|
||
def _fstart_from_geometries(rot_img, t_img):
|
||
"""從旋轉檔與未旋轉模板檔的 origin 差回復 fstart（模板 index 系，
|
||
(x,y,z) 整數）；幾何不一致（非對角 direction、spacing / direction
|
||
不同）或差不是整數時回傳 None。"""
|
||
ax_r, ax_t = _axis_vectors(rot_img), _axis_vectors(t_img)
|
||
if ax_r is None or ax_t is None:
|
||
return None
|
||
(d_r, s_r), (d_t, s_t) = ax_r, ax_t
|
||
if not (np.allclose(d_r, d_t, atol=1e-6) and np.allclose(s_r, s_t, atol=1e-6)):
|
||
return None
|
||
o_r = np.asarray(rot_img.GetOrigin(), dtype=float)
|
||
o_t = np.asarray(t_img.GetOrigin(), dtype=float)
|
||
fs = (o_r - o_t) / (d_t * s_t)
|
||
if not np.allclose(fs, np.round(fs), atol=1e-6):
|
||
return None
|
||
return np.round(fs).astype(int)
|
||
|
||
|
||
def _std_to_orig_affine(orig_img, t_img, ap_flip):
|
||
"""std（0.5mm 裁切檔、standardize_affine 後、ap_flip 後）-> 原始 CT
|
||
物理座標（LPS）：P_orig = A·P_std + b。回傳 (A (3,) 對角, b (3,))；
|
||
原始 direction 非純對角時回傳 None。t_img = 該 level 的未旋轉標準化
|
||
裁切檔（disk 上的後 standardize 幾何）。
|
||
|
||
standardize_affine 是【逐檔（每 level 裁切）】處理：翻軸（nibabel：
|
||
affine 對角與平移同反號 + 資料 flip）在 worlds 中是整份資料平移
|
||
t_i = 2·O_post_i - d_i·h·(L_i-1)（O_post/L 為 disk 上該檔的
|
||
origin/尺寸、d 為原始 direction、h=0.5；翻軸條件 RAS 對角 < 0
|
||
= x: d_x=+1、y: d_y=+1、z: d_z=-1），故 b_i = -t_i（翻軸的軸）。
|
||
ap_flip 是 y 軸純資料鏡射（幾何不變，作用在整張 0.5mm full grid）：
|
||
y_ap = C - y_orig，C = 2·O_full_y + d_y·h·(L_full_y-1)（full grid
|
||
幾何 = 原始 CT），發生在 standardize 之前。"""
|
||
d_full = np.asarray(orig_img.GetDirection(), dtype=float).reshape(3, 3)
|
||
d = np.diag(d_full)
|
||
if not (np.allclose(d_full, np.diag(d), atol=1e-6)
|
||
and np.allclose(np.abs(d), 1.0, atol=1e-6)):
|
||
return None
|
||
h = 0.5
|
||
O_full = np.asarray(orig_img.GetOrigin(), dtype=float)
|
||
s = np.asarray(orig_img.GetSpacing(), dtype=float)
|
||
n = np.asarray(orig_img.GetSize(), dtype=float) # (x, y, z)
|
||
L_full = np.ceil(n * s / h - 1e-6) # 0.5mm full grid
|
||
O_post = np.asarray(t_img.GetOrigin(), dtype=float)
|
||
L = np.asarray(t_img.GetSize(), dtype=float) # 裁切檔尺寸
|
||
flips = (d[0] == 1.0, d[1] == 1.0, d[2] == -1.0)
|
||
A = np.ones(3)
|
||
b = np.zeros(3)
|
||
for i in range(3):
|
||
if flips[i]:
|
||
b[i] = -(2.0 * O_post[i] - d[i] * h * (L[i] - 1.0))
|
||
if ap_flip:
|
||
# y0 --AP 鏡射--> y1 = C - y0 --std 平移--> y_std = y1 + t_y
|
||
# => y0 = -y_std + (C + t_y)
|
||
A[1] = -1.0
|
||
C_full = 2.0 * O_full[1] + d[1] * h * (L_full[1] - 1.0)
|
||
if flips[1]:
|
||
t_y = 2.0 * O_post[1] - d[1] * h * (L[1] - 1.0)
|
||
b[1] = C_full + t_y
|
||
else:
|
||
b[1] = C_full
|
||
return A, b
|
||
|
||
|
||
def _affine_sitk(A, b):
|
||
"""對角 A + 平移 b 的 sitk 仿射（LPS 物理 -> LPS 物理）。
|
||
p' = A·p + b；SetMatrix 為 row-major 3x3、SetTranslation 獨立。
|
||
自檢三軸 + 原點。"""
|
||
tr = sitk.AffineTransform(3)
|
||
tr.SetMatrix([float(A[0]), 0.0, 0.0,
|
||
0.0, float(A[1]), 0.0,
|
||
0.0, 0.0, float(A[2])])
|
||
tr.SetTranslation([float(b[0]), float(b[1]), float(b[2])])
|
||
for i in range(3):
|
||
p = [0.0, 0.0, 0.0]
|
||
p[i] = 123.0
|
||
got = np.array(tr.TransformPoint(tuple(p)))
|
||
exp = np.array(b, dtype=float)
|
||
exp[i] = A[i] * 123.0 + b[i]
|
||
if not np.allclose(got, exp, atol=1e-6):
|
||
raise RuntimeError(f'affine self-check failed: axis {i} '
|
||
f'{got} != {exp}')
|
||
return tr
|
||
|
||
|
||
def _rotate_back(arr, R, c_xyz, fstart, shape_zyx, order=0, cval=0.0):
|
||
"""rotated grid 上的 field 反向旋轉到未旋轉模板 grid（shape_zyx）。
|
||
rotated 檔 array 座標 a = 模板座標 V - fstart；反向映射
|
||
V = R(W - c) + c，W 為模板 grid 上每格（x,y,z 向量，同
|
||
compute_normalizing_rotation 慣例）。order / cval 同 map_coordinates。"""
|
||
nz, ny, nx = shape_zyx
|
||
idx = np.indices((nz, ny, nx), dtype=np.float64)
|
||
W = np.stack([idx[2], idx[1], idx[0]], axis=0) # (3, z, y, x)
|
||
c = np.asarray(c_xyz, dtype=np.float64)[:, None, None, None]
|
||
Rf = np.asarray(R, dtype=np.float64)
|
||
V = np.tensordot(Rf, W - c, axes=([1], [0])) + c
|
||
fs = np.asarray(fstart, dtype=np.float64)[:, None, None, None]
|
||
a = V - fs # rotated 檔 array 座標
|
||
return map_coordinates(arr, [a[2], a[1], a[0]],
|
||
order=order, cval=cval, mode='constant')
|
||
|
||
|
||
# ---------------------------------------------------------------- level 處理
|
||
def _load_unrotated_template(vol_dir, level):
|
||
"""_write_rotated_level 的旋轉來源：_smd_resampled（有效時）否則
|
||
_binary_nn。回傳 (img, arr, is_smd)；皆無則 None。
|
||
未旋轉檔位置經 level_file_path（新世代 crop/、舊世代頂層）。"""
|
||
smd_path = level_file_path(vol_dir, level, 'smd_resampled')
|
||
if os.path.exists(smd_path):
|
||
sa = sitk.GetArrayFromImage(sitk.ReadImage(smd_path)).astype(np.float32)
|
||
if sa.size > 0 and np.isfinite(sa).all() and sa.min() < 0.0 and sa.max() > 0.0:
|
||
return sitk.ReadImage(smd_path), sa, True
|
||
nn_path = level_file_path(vol_dir, level, 'binary_nn')
|
||
if os.path.exists(nn_path):
|
||
na = sitk.GetArrayFromImage(sitk.ReadImage(nn_path))
|
||
if na.size > 0 and int((na > 0).sum()) > 0:
|
||
return sitk.ReadImage(nn_path), na, False
|
||
return None
|
||
|
||
|
||
def _bone_mask(t_arr, is_smd):
|
||
return (t_arr < 0.5) if is_smd else (t_arr > 0)
|
||
|
||
|
||
def _forward_identity_check(vol_dir, level, R, c_xyz, t_img, t_arr, is_smd):
|
||
"""以 _write_rotated_level 完全同參數（fstart/fsize 由未旋轉 _binary_nn
|
||
旋轉後 bbox±4、order / cval 同 writer）重旋轉、與存檔的旋轉檔逐體素
|
||
比對，回傳 mismatch 比例（(R,c) 相同時 ≈0，精確驗證、不受物件厚度
|
||
影響）；無可比檔（缺檔 / 形狀不符）時回 None。"""
|
||
rotated_dir = os.path.join(vol_dir, 'rotated')
|
||
if is_smd:
|
||
target = os.path.join(rotated_dir, f'{level}_smd_resampled.nii.gz')
|
||
src_arr, order, cval = t_arr, 1, float(t_arr.max())
|
||
else:
|
||
target = os.path.join(rotated_dir, f'{level}_binary_nn.nii.gz')
|
||
src_arr, order, cval = (t_arr > 0).astype(np.uint8), 0, 0.0
|
||
if not os.path.exists(target):
|
||
return None
|
||
start, size = rotated_grid(t_arr.shape, R, c_xyz, margin=4)
|
||
fstart = (int(start[0]), int(start[1]), int(start[2]))
|
||
fsize = (int(size[0]), int(size[1]), int(size[2]))
|
||
nn_path = level_file_path(vol_dir, level, 'binary_nn')
|
||
if os.path.exists(nn_path):
|
||
nn_img = sitk.ReadImage(nn_path)
|
||
nn_arr = (sitk.GetArrayFromImage(
|
||
sitk.Resample(nn_img, t_img, interpolator=sitk.sitkNearestNeighbor,
|
||
defaultPixelValue=0)) > 0).astype(np.uint8)
|
||
rot_nn_full = (rotate_volume_to(nn_arr, R, c_xyz, fstart, fsize,
|
||
order=0, cval=0.0) > 0.5).astype(np.uint8)
|
||
if int(rot_nn_full.sum()) > 0:
|
||
zz, yy, xx = np.where(rot_nn_full > 0)
|
||
m = 4
|
||
z0 = max(0, int(zz.min()) - m)
|
||
z1 = min(rot_nn_full.shape[0] - 1, int(zz.max()) + m)
|
||
y0 = max(0, int(yy.min()) - m)
|
||
y1 = min(rot_nn_full.shape[1] - 1, int(yy.max()) + m)
|
||
x0 = max(0, int(xx.min()) - m)
|
||
x1 = min(rot_nn_full.shape[2] - 1, int(xx.max()) + m)
|
||
fstart = (fstart[0] + x0, fstart[1] + y0, fstart[2] + z0)
|
||
fsize = (x1 - x0 + 1, y1 - y0 + 1, z1 - z0 + 1)
|
||
out = rotate_volume_to(src_arr, R, c_xyz, fstart, fsize, order=order,
|
||
cval=cval)
|
||
saved = sitk.GetArrayFromImage(sitk.ReadImage(target))
|
||
if out.shape != saved.shape:
|
||
return None
|
||
if order == 0:
|
||
return float(((out > 0.5) != (saved > 0.5)).mean())
|
||
return float((np.abs(out.astype(np.float64)
|
||
- saved.astype(np.float64)) > 0.02).mean())
|
||
|
||
|
||
def _containment_roundtrip(vol_dir, level, R, c_xyz, fstart, t_arr, is_smd,
|
||
it=2):
|
||
"""fallback 驗證：旋轉 check 檔反向旋轉後與未旋轉 bone mask 互相包含於
|
||
it-voxel 膨脹內（薄結構的 round-trip IoU 天生低，用包含率代替）。
|
||
回傳 min 方向包含率；無 check 檔時回 None。"""
|
||
from scipy.ndimage import binary_dilation
|
||
rotated_dir = os.path.join(vol_dir, 'rotated')
|
||
candidates = [
|
||
(f'{level}_smd_resampled.nii.gz', 1),
|
||
(f'{level}_binary_sdf.nii.gz', 0),
|
||
(f'{level}_binary_nn.nii.gz', 0),
|
||
]
|
||
if not is_smd:
|
||
candidates = candidates[1:]
|
||
B0 = _bone_mask(t_arr, is_smd)
|
||
for fname, order in candidates:
|
||
p = os.path.join(rotated_dir, fname)
|
||
if not os.path.exists(p):
|
||
continue
|
||
a = sitk.GetArrayFromImage(sitk.ReadImage(p)).astype(np.float64)
|
||
if a.size == 0:
|
||
continue
|
||
cval = float(a.max()) if fname.endswith('_smd_resampled.nii.gz') else 0.0
|
||
b = _rotate_back(a, R, c_xyz, fstart, t_arr.shape, order=order,
|
||
cval=cval)
|
||
B1 = b < 0.5 if fname.endswith('_smd_resampled.nii.gz') else b > 0.5
|
||
if not B1.any() or not B0.any():
|
||
continue
|
||
da = binary_dilation(B1, iterations=it)
|
||
db = binary_dilation(B0, iterations=it)
|
||
return float(min((B1 & db).sum() / B1.sum(), (B0 & da).sum() / B0.sum()))
|
||
return None
|
||
|
||
|
||
def process_level(vol_dir, level, min_contain):
|
||
"""單一 level：重算 (R, c)、回復 fstart、驗證、反向旋轉 label。
|
||
回傳 (unrotated_label_arr (z,y,x) uint8 0/1/2/3, t_img, info)；失敗回 None。"""
|
||
name = os.path.basename(vol_dir)
|
||
rot_lbl_path = os.path.join(vol_dir, 'rotated', f'{level}_label.nii.gz')
|
||
if not os.path.exists(rot_lbl_path):
|
||
return None
|
||
rot_img = sitk.ReadImage(rot_lbl_path)
|
||
rot_arr = sitk.GetArrayFromImage(rot_img).astype(np.uint8)
|
||
|
||
tpl = _load_unrotated_template(vol_dir, level)
|
||
if tpl is None:
|
||
logger.warning(f'{name} {level}: 無可用未旋轉模板 '
|
||
f'（_smd_resampled / _binary_nn），跳過')
|
||
return None
|
||
t_img, t_arr, is_smd = tpl
|
||
|
||
fstart = _fstart_from_geometries(rot_img, t_img)
|
||
if fstart is None:
|
||
logger.warning(f'{name} {level}: 旋轉 / 未旋轉幾何不一致，跳過')
|
||
return None
|
||
|
||
bin_arr = _bone_mask(t_arr, is_smd).astype(np.uint8)
|
||
sym = best_symmetry_plane(bin_arr)
|
||
symp = best_upper_endplate_plane(bin_arr)
|
||
R, c_xyz = compute_normalizing_rotation(bin_arr, sym, symp)
|
||
|
||
# 驗證重算的 (R, c)：主要用 forward identity（與生成時同參數重旋轉、
|
||
# 逐體素比對存檔；(R,c) 相同時 ≈0）；無可比檔時 fallback 用
|
||
# round-trip 包含率（min_contain 門檻）。
|
||
fwd = _forward_identity_check(vol_dir, level, R, c_xyz, t_img, t_arr, is_smd)
|
||
contain = None
|
||
if fwd is not None:
|
||
if fwd > 1e-4:
|
||
logger.warning(f'{name} {level}: 正向重旋轉 mismatch {fwd:.4f} '
|
||
f'> 1e-4，(R,c) 與生成參數不符，跳過')
|
||
return None
|
||
else:
|
||
contain = _containment_roundtrip(vol_dir, level, R, c_xyz, fstart,
|
||
t_arr, is_smd)
|
||
if contain is None:
|
||
logger.warning(f'{name} {level}: 無 check 檔可驗證，跳過')
|
||
return None
|
||
if contain < min_contain:
|
||
logger.warning(f'{name} {level}: round-trip 包含率 {contain:.4f} '
|
||
f'< {min_contain}，跳過該 level')
|
||
return None
|
||
|
||
unrot = _rotate_back(rot_arr.astype(np.float64), R, c_xyz, fstart,
|
||
t_arr.shape).astype(np.uint8)
|
||
n_lbl = int((unrot > 0).sum())
|
||
B0 = _bone_mask(t_arr, is_smd)
|
||
bone_frac = (int(((unrot > 0) & B0).sum()) / n_lbl) if n_lbl else 0.0
|
||
info = {
|
||
'forward_mismatch': fwd,
|
||
'containment': contain,
|
||
'fstart': [int(v) for v in fstart],
|
||
'bone_frac': bone_frac,
|
||
'rotated': {'vbody': int((rot_arr == 1).sum()),
|
||
'sp': int((rot_arr == 2).sum()),
|
||
'other': int((rot_arr == 3).sum())},
|
||
'unrotated': {'vbody': int((unrot == 1).sum()),
|
||
'sp': int((unrot == 2).sum()),
|
||
'other': int((unrot == 3).sum())},
|
||
}
|
||
return unrot, t_img, info
|
||
|
||
|
||
def _resample_nearest(ref_img, src_img, transform=None):
|
||
r = sitk.ResampleImageFilter()
|
||
r.SetReferenceImage(ref_img)
|
||
r.SetInterpolator(sitk.sitkNearestNeighbor)
|
||
r.SetDefaultPixelValue(0)
|
||
if transform is not None:
|
||
r.SetTransform(transform)
|
||
return r.Execute(src_img)
|
||
|
||
|
||
def process_volume(vol_dir, orig_img, key, dest_dir, min_contain, force,
|
||
ap_flip):
|
||
"""所有 level 反旋轉 + 合併 + 重取樣到原始 CT grid，存「一個」.nii.gz。
|
||
回傳 (out_path, status, {level: info}, vol_info)；失敗時 out_path=None、
|
||
status 為原因。"""
|
||
name = os.path.basename(vol_dir)
|
||
nxi, nyi, nzi = orig_img.GetSize() # sitk GetSize = (x, y, z)
|
||
merged = np.zeros((nzi, nyi, nxi), dtype=np.uint8) # array = (z, y, x)
|
||
level_info = {}
|
||
maps = []
|
||
|
||
for li, level in enumerate(LUMBAR_LEVELS):
|
||
res = process_level(vol_dir, level, min_contain)
|
||
if res is None:
|
||
continue
|
||
unrot, t_img, info = res
|
||
# std->orig mapping 依該 level 裁切檔幾何逐 level 算
|
||
aff = _std_to_orig_affine(orig_img, t_img, bool(ap_flip))
|
||
if aff is None:
|
||
logger.warning(f'{name} {level}: 原始 CT direction 非純對角，'
|
||
f'跳過')
|
||
continue
|
||
A, b = aff
|
||
# Resample 的 transform 方向是 dest(原始) physical -> src(std)
|
||
# physical：P_std = A·(P_orig - b) = A·P_orig - A*b（A 對角 ±1）
|
||
tr = _affine_sitk(A, -A * b)
|
||
# 重編號：Lx VBODY=2x-1、SP=2x（x = L1..L6 的 1-based）；
|
||
# other(3) 不納入
|
||
new = np.zeros_like(unrot)
|
||
new[unrot == 1] = 2 * li + 1
|
||
new[unrot == 2] = 2 * li + 2
|
||
maps.append({'A': [float(v) for v in A],
|
||
'b': [round(float(v), 3) for v in b]})
|
||
if not new.any():
|
||
logger.info(f'{name} {level}: 無 VBODY / SP voxel，不計入')
|
||
continue
|
||
lvl_img = sitk.GetImageFromArray(new)
|
||
lvl_img.CopyInformation(t_img)
|
||
arr = sitk.GetArrayFromImage(_resample_nearest(orig_img, lvl_img, tr))
|
||
merged = np.where(arr > 0, arr, merged)
|
||
info['orig'] = {'vbody': int((arr == 2 * li + 1).sum()),
|
||
'sp': int((arr == 2 * li + 2).sum())}
|
||
level_info[level] = info
|
||
check = (f'fwd_mismatch={info["forward_mismatch"]:.1e}'
|
||
if info.get('forward_mismatch') is not None
|
||
else f'containment={info["containment"]:.4f}')
|
||
logger.info(f'{name} {level}: {check} '
|
||
f'bone_frac={info["bone_frac"]:.4f} '
|
||
f'std->orig A={np.round(A, 3)} b={np.round(b, 2)} '
|
||
f'orig vbody={info["orig"]["vbody"]} '
|
||
f'sp={info["orig"]["sp"]} (labels {2 * li + 1}/{2 * li + 2})')
|
||
|
||
if not level_info:
|
||
return None, 'no usable level', {}, {}
|
||
if int(merged.sum()) == 0:
|
||
return None, 'merged label empty after resample', level_info, {}
|
||
|
||
# 對位檢查：label 應落在原始 CT 體內（HU > -100）且主要是骨頭
|
||
#（median HU >= 100；鬆質骨 HU 可 < 100，故不用 bone 比例作硬門檻）
|
||
ct = sitk.GetArrayFromImage(orig_img)
|
||
hu = ct[merged > 0]
|
||
frac_body = float((hu > -100).mean())
|
||
med_hu = float(np.median(hu))
|
||
frac_bone = float((hu > 100).mean())
|
||
vol_info = {'in_body_frac': frac_body, 'bone_frac': frac_bone,
|
||
'median_hu': med_hu, 'std_to_orig': maps, 'ap_flip': bool(ap_flip)}
|
||
if frac_body < 0.98 or med_hu < 100:
|
||
logger.warning(f'{name}: 對位檢查 in_body={frac_body:.4f} '
|
||
f'median HU={med_hu:.0f}，疑似錯位，跳過')
|
||
return None, f'alignment check failed ' \
|
||
f'(in_body={frac_body:.3f}, median_hu={med_hu:.0f})', \
|
||
level_info, vol_info
|
||
logger.info(f'{name}: alignment OK in_body={frac_body:.4f} '
|
||
f'bone_frac={frac_bone:.4f} median HU={med_hu:.0f}')
|
||
|
||
out_path = os.path.join(dest_dir, key, f'{name}.nii.gz')
|
||
if os.path.exists(out_path) and not force:
|
||
return out_path, 'exists (skipped)', level_info, vol_info
|
||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||
out_img = sitk.GetImageFromArray(merged)
|
||
out_img.CopyInformation(orig_img)
|
||
sitk.WriteImage(out_img, out_path)
|
||
logger.info(f'{name}: saved {out_path} (levels={list(level_info)})')
|
||
return out_path, 'written', level_info, vol_info
|
||
|
||
|
||
def find_original_ct(name):
|
||
"""在 DATA_ROOT 子目錄裡找 <name>.nii.gz；回傳 (path, 子目錄) 或
|
||
(None, None)。"""
|
||
if not os.path.isdir(DATA_ROOT):
|
||
return None, None
|
||
for key in sorted(os.listdir(DATA_ROOT)):
|
||
p = os.path.join(DATA_ROOT, key, f'{name}.nii.gz')
|
||
if os.path.isfile(p):
|
||
return p, key
|
||
return None, None
|
||
|
||
|
||
def list_volumes_with_labels(root):
|
||
"""root 下有 rotated/*_label.nii.gz 的 volume 目錄（依名稱排序）。"""
|
||
vols = []
|
||
if not os.path.isdir(root):
|
||
return vols
|
||
for d in sorted(os.listdir(root)):
|
||
rot = os.path.join(root, d, 'rotated')
|
||
if os.path.isdir(rot) and any(
|
||
f.endswith('_label.nii.gz') for f in os.listdir(rot)):
|
||
vols.append(os.path.join(root, d))
|
||
return vols
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='Merge unrotated lumbar VBODY/SP labels into one '
|
||
'original-CT coordinate .nii.gz per patient.')
|
||
parser.add_argument('--roots', nargs='*', default=DEFAULT_ROOTS,
|
||
help='Output generation dir(s) to scan (newest first).')
|
||
parser.add_argument('--dest', default=DEFAULT_DEST,
|
||
help=f'Destination dir (default: {DEFAULT_DEST})')
|
||
parser.add_argument('--force', action='store_true',
|
||
help='Overwrite existing outputs.')
|
||
parser.add_argument('--min-contain', type=float, default=0.85,
|
||
help='Fallback per-level check: skip a level whose '
|
||
'unrotation round-trip containment is below '
|
||
'this (default: 0.85).')
|
||
parser.add_argument('--name', default=None,
|
||
help='Process only this volume name (debug).')
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(LOG_DIR, exist_ok=True)
|
||
log_path = os.path.join(LOG_DIR, f'xfr_orig_labels_{time.strftime("%Y%m%d_%H%M%S")}.log')
|
||
setup_tee(log_path)
|
||
logging.basicConfig(level=logging.INFO,
|
||
format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
|
||
datefmt='%Y-%m-%d %H:%M:%S')
|
||
logger.info(f'Log file: {log_path}')
|
||
logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}')
|
||
logger.info(f'Roots: {args.roots}')
|
||
logger.info(f'Dest: {args.dest}')
|
||
logger.info('Label scheme: Lx VBODY=2x-1, SP=2x (L1:1/2, L2:3/4, '
|
||
'L3:5/6, L4:7/8, L5:9/10, L6:11/12); other bone excluded; '
|
||
'uint8; geometry = original CT')
|
||
|
||
os.makedirs(args.dest, exist_ok=True)
|
||
ap_flip_db = load_ap_flip()
|
||
logger.info(f'ap_flip metadata: {len(ap_flip_db)} entries '
|
||
f'({sum(1 for v in ap_flip_db.values() if v)} flipped)')
|
||
done, summary, no_ap = set(), {}, set()
|
||
summary_path = os.path.join(args.dest, 'summary.json')
|
||
if os.path.exists(summary_path):
|
||
try:
|
||
with open(summary_path) as f:
|
||
summary = json.load(f)
|
||
except (json.JSONDecodeError, OSError):
|
||
summary = {}
|
||
|
||
for root in args.roots:
|
||
vols = list_volumes_with_labels(root)
|
||
if args.name is not None:
|
||
vols = [v for v in vols if os.path.basename(v) == args.name]
|
||
logger.info(f'Scanning {root}: {len(vols)} volume(s) with rotated labels')
|
||
for vol_dir in vols:
|
||
name = os.path.basename(vol_dir)
|
||
if name in done:
|
||
logger.info(f'[skip] {name}: 已由較新世代處理')
|
||
continue
|
||
done.add(name)
|
||
|
||
orig_path, key = find_original_ct(name)
|
||
if orig_path is None:
|
||
logger.warning(f'[skip] {name}: 找不到原始 CT '
|
||
f'({DATA_ROOT}*/{name}.nii.gz)')
|
||
continue
|
||
ap_flip = ap_flip_db.get(name)
|
||
if ap_flip is None and name not in no_ap:
|
||
no_ap.add(name)
|
||
logger.warning(f'{name}: metadata db 缺 ap_flip，假設 False '
|
||
f'（若為 AP 翻轉個案將對位失敗並被抓出）')
|
||
ap_flip = False
|
||
out_path = os.path.join(args.dest, key, f'{name}.nii.gz')
|
||
if os.path.exists(out_path) and not args.force:
|
||
logger.info(f'[skip] {name}: {out_path} 已存在')
|
||
continue
|
||
|
||
try:
|
||
orig_img = sitk.ReadImage(orig_path)
|
||
t0 = time.time()
|
||
out, status, level_info, vol_info = process_volume(
|
||
vol_dir, orig_img, key, args.dest, args.min_iou, args.force,
|
||
ap_flip)
|
||
dt = time.time() - t0
|
||
common = {'status': status, 'source_root': root,
|
||
'vol_info': vol_info,
|
||
'updated': time.strftime('%Y-%m-%d %H:%M:%S')}
|
||
if out is None:
|
||
logger.warning(f'[skip] {name}: {status} ({dt:.1f}s)')
|
||
summary[name] = {**common, 'dest_root': DATA_ROOT,
|
||
'levels': level_info}
|
||
else:
|
||
logger.info(f'[done] {name}: {status} in {dt:.1f}s')
|
||
summary[name] = {**common, 'path': out, 'key': key,
|
||
'levels': level_info}
|
||
except Exception as e:
|
||
logger.warning(f'[error] {name}: {e}')
|
||
summary[name] = {'status': f'error: {e}', 'source_root': root,
|
||
'updated': time.strftime('%Y-%m-%d %H:%M:%S')}
|
||
|
||
with open(summary_path, 'w') as f:
|
||
json.dump(summary, f, indent=2, ensure_ascii=False)
|
||
logger.info('All done.')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |