Xiao Furen
264e17b13a
- Implemented xfr_pso.py to handle PSO optimization tasks for lumbar levels, including GPU management and logging. - Introduced xfr_record_planes.py to supplement existing volumes with symmetry and endplate plane data in transform.json. - Added functionality to skip already recorded levels unless forced, ensuring efficient processing. - Enhanced logging for better traceability of operations and errors during execution.
214 lines
No EOL
8.3 KiB
Python
214 lines
No EOL
8.3 KiB
Python
#!/home/xfr/.conda/envs/cbt/bin/python
|
||
"""補記既有 volume 的 transform.json rotated 段平面（sym / endplate）。
|
||
|
||
core/optimizer 的 xfr 管線（run_pso_torch_xfr / run_de_torch_xfr）不再於影像
|
||
上重新估算兩平面（best_symmetry_plane / best_upper_endplate_plane），改直接
|
||
讀 <vol>/transform.json levels[L*]['rotated']['sym' / 'endplate']
|
||
（imaging.transforms.load_rotated_planes_from_path）：
|
||
- 新世代：xfr_preprocess._write_rotated_level 旋轉時即記錄（rotated/ 局部系）；
|
||
- 舊世代（進度 progress.json 已標完成、重跑 preprocess 會整支跳過）：跑本腳本
|
||
補記——用與舊 optimizer 相同函式、相同預設值，對
|
||
rotated/<level>_binary_sdf.nii.gz（完整 mask、SP 移除前）重新估算，
|
||
併入既有 rotated 段（只加 sym / endplate 兩鍵，不動 template / R / center /
|
||
start / size 等既有鍵、不重寫任何 nii）。補記後 optimizer 讀到的平面與
|
||
改前行為一致（同輸入、同演算法）。
|
||
|
||
預設跳過已有 sym 記錄的 level（部分失敗後重跑只補缺者）；--force 全數重記。
|
||
|
||
Usage:
|
||
python xfr_record_planes.py [--force] [--dir <standardized dir>] [volume_id ...]
|
||
volume_id 可用完整 ID 或末段（如 0001）；省略 = 全部 volume。
|
||
"""
|
||
import argparse
|
||
import logging
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
|
||
from imaging.orientation import best_symmetry_plane, best_upper_endplate_plane
|
||
from imaging.transforms import load_transform, save_transform
|
||
|
||
|
||
STANDARDIZED_DIR = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'
|
||
|
||
LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
|
||
logger = logging.getLogger('xfr_record_planes')
|
||
|
||
# 基準系健全閾值：對齊良好時 theta_v / tau_y 應近似 0（xfr_preprocess 的 R
|
||
# 已把鏡稱面法線轉到 +x、終板法線轉到 z>0）
|
||
SANITY_DEG = 5.0
|
||
|
||
|
||
class _Tee:
|
||
"""同時輸出到 console 與 log 檔，逐行寫入。"""
|
||
|
||
def __init__(self, console, log_fh):
|
||
self.console = console
|
||
self.log_fh = log_fh
|
||
self.buf = ''
|
||
|
||
def _emit(self, line):
|
||
self.console.write(line + '\n')
|
||
self.log_fh.write(line + '\n')
|
||
|
||
def write(self, data):
|
||
if not data:
|
||
return
|
||
self.buf += data
|
||
while True:
|
||
idx_n = self.buf.find('\n')
|
||
idx_r = self.buf.find('\r')
|
||
candidates = [i for i in (idx_n, idx_r) if i != -1]
|
||
if not candidates:
|
||
break
|
||
idx = min(candidates)
|
||
self._emit(self.buf[:idx])
|
||
self.buf = self.buf[idx + 1:]
|
||
|
||
def flush(self):
|
||
if self.buf:
|
||
self._emit(self.buf)
|
||
self.buf = ''
|
||
self.console.flush()
|
||
self.log_fh.flush()
|
||
|
||
def isatty(self):
|
||
return False
|
||
|
||
|
||
def setup_tee(log_path):
|
||
log_fh = open(log_path, 'a', buffering=1)
|
||
sys.stdout = _Tee(sys.stdout, log_fh)
|
||
sys.stderr = _Tee(sys.stderr, log_fh)
|
||
|
||
|
||
def rotated_levels(vol_dir):
|
||
"""該 volume rotated/ 中有 _binary_sdf.nii.gz 的 lumbar level（排序）。"""
|
||
rotated_dir = os.path.join(vol_dir, 'rotated')
|
||
if not os.path.isdir(rotated_dir):
|
||
return ()
|
||
levels = set()
|
||
for fn in os.listdir(rotated_dir):
|
||
m = re.fullmatch(r'(L[1-9]\d*)_binary_sdf\.nii\.gz', fn)
|
||
if m:
|
||
levels.add(m.group(1))
|
||
return tuple(sorted(levels, key=lambda lv: int(lv[1:])))
|
||
|
||
|
||
def record_level(vol_dir, level, sdf_path):
|
||
"""估算該 level 兩平面併入 transform.json rotated 段；回傳 (sym, endplate)。"""
|
||
arr = sitk.GetArrayFromImage(sitk.ReadImage(sdf_path))
|
||
sym = best_symmetry_plane(arr)
|
||
endplate = best_upper_endplate_plane(arr)
|
||
|
||
meta = load_transform(vol_dir)
|
||
lv = meta['levels'].setdefault(level, {})
|
||
rot = lv.setdefault('rotated', {})
|
||
rot['sym'] = sym
|
||
rot['endplate'] = endplate
|
||
save_transform(vol_dir, meta)
|
||
|
||
s_nx, s_ny, _ = sym['normal']
|
||
theta_v = float(np.degrees(np.arctan2(s_ny, s_nx)))
|
||
if endplate is not None:
|
||
e_nx, e_ny, e_nz = endplate['normal']
|
||
tau_y = float(np.degrees(np.arctan2(e_ny, e_nz)))
|
||
else:
|
||
tau_y = None
|
||
note = f'endplate={endplate is not None}'
|
||
if abs(theta_v) > SANITY_DEG or (tau_y is not None and abs(tau_y) > SANITY_DEG):
|
||
logger.warning(f'{os.path.basename(vol_dir)} {level}: 基準系殘差偏大 '
|
||
f'(theta_v={theta_v:+.2f} deg, tau_y='
|
||
f'{tau_y if tau_y is None else round(tau_y, 2)}) — '
|
||
f'該 level 的 xfr_preprocess 對齊可能失效，建議人工核對')
|
||
else:
|
||
logger.info(f'{os.path.basename(vol_dir)} {level}: recorded '
|
||
f'(mirror ratio={sym["ratio"]:.3f}, theta_v={theta_v:+.2f} deg, '
|
||
f'tau_y={tau_y if tau_y is None else round(tau_y, 2)}, {note})')
|
||
return sym, endplate
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='Record sym/endplate planes into transform.json for '
|
||
'existing volumes (old generation, before xfr_preprocess '
|
||
'recording them at rotation time).')
|
||
parser.add_argument('volume_ids', nargs='*',
|
||
help='Full ID or trailing segment (e.g. 0001); '
|
||
'default: all volumes')
|
||
parser.add_argument('--dir', default=STANDARDIZED_DIR,
|
||
help='Standardized output dir (default: standardized-xfr-3)')
|
||
parser.add_argument('--force', action='store_true',
|
||
help='Re-record levels that already have sym/endplate')
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(LOG_DIR, exist_ok=True)
|
||
log_path = os.path.join(
|
||
LOG_DIR, f'xfr_record_planes_{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)}')
|
||
|
||
if not os.path.isdir(args.dir):
|
||
sys.exit(f'Standardized dir not found: {args.dir}')
|
||
volumes = [d for d in sorted(os.listdir(args.dir))
|
||
if os.path.isdir(os.path.join(args.dir, d))]
|
||
if args.volume_ids:
|
||
keys = [v.lower() for v in args.volume_ids]
|
||
vols = [v for v in volumes
|
||
if v.lower() in keys or v.rsplit('.', 1)[-1].lower() in keys]
|
||
if not vols:
|
||
sys.exit(f'No volume matches: {", ".join(args.volume_ids)}')
|
||
volumes = vols
|
||
logger.info(f'{len(volumes)} volume(s) under {args.dir}'
|
||
+ (' (force)' if args.force else ''))
|
||
|
||
t0 = time.time()
|
||
n_rec, n_skip_done, n_err, n_no_levels = 0, 0, 0, 0
|
||
no_ep = []
|
||
for vid in volumes:
|
||
vol_dir = os.path.join(args.dir, vid)
|
||
levels = rotated_levels(vol_dir)
|
||
if not levels:
|
||
n_no_levels += 1
|
||
continue
|
||
try:
|
||
meta = load_transform(vol_dir)
|
||
except FileNotFoundError:
|
||
logger.warning(f'{vid}: transform.json 缺失，跳過 '
|
||
f'（舊世代可先在 xfr_inverse_transform.py 加 --rebuild）')
|
||
n_err += 1
|
||
continue
|
||
for level in levels:
|
||
rot = meta['levels'].get(level, {}).get('rotated') or {}
|
||
if not args.force and rot.get('sym') is not None:
|
||
n_skip_done += 1
|
||
continue
|
||
sdf_path = os.path.join(vol_dir, 'rotated', f'{level}_binary_sdf.nii.gz')
|
||
try:
|
||
_, ep = record_level(vol_dir, level, sdf_path)
|
||
n_rec += 1
|
||
if ep is None:
|
||
no_ep.append(f'{vid} {level}')
|
||
except Exception as e:
|
||
n_err += 1
|
||
logger.error(f'{vid} {level}: {e}')
|
||
dt = time.time() - t0
|
||
logger.info(f'Done in {dt / 60:.1f} min: recorded {n_rec} level(s), '
|
||
f'skipped(already) {n_skip_done}, no rotated levels '
|
||
f'{n_no_levels} volume(s), errors {n_err}')
|
||
if no_ep:
|
||
logger.warning(f'{len(no_ep)} level(s) 無上終板面記錄（資料不足），'
|
||
f'optimizer 對應 level 終板入口剪除停用: {", ".join(no_ep[:10])}'
|
||
f'{" ..." if len(no_ep) > 10 else ""}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |