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.
466 lines
No EOL
17 KiB
Python
466 lines
No EOL
17 KiB
Python
#!/home/xfr/.conda/envs/cbt/bin/python
|
||
|
||
import logging
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
import queue as queue_module
|
||
import subprocess
|
||
import multiprocessing as mp
|
||
|
||
import SimpleITK as sitk
|
||
import torch
|
||
|
||
# from config.device import get_device
|
||
from core.cylinder import create_coordinate_grid
|
||
from core.objective import set_global_context
|
||
from core.optimizer import run_pso_torch, run_de_torch, run_nm_torch, run_pso_torch_xfr
|
||
from imaging.orientation import azimuth_rotation, analyze_vertebral_tilt_contour
|
||
from imaging.transforms import level_file_path
|
||
|
||
import xfr_cbt_native
|
||
|
||
standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'
|
||
|
||
azimuth_rotation_dir = '/mnt/1248/open2/cyrou/azimuth_rotation'
|
||
tilt_contour_dir = '/mnt/1248/open2/cyrou/tilt_contour'
|
||
Output_dir = '/mnt/1248/open/cyrou/Output'
|
||
|
||
LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')
|
||
|
||
LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
|
||
|
||
# 單側任務的中繼結果（<run_id>/<volume_id>/<level>_<side>.json + plot lock）：
|
||
# 同一 level 的 L/R 可跑在不同 GPU，較晚完成的一側讀到兩側結果後跑合併輸出
|
||
SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results')
|
||
|
||
logger = logging.getLogger('xfr_debug')
|
||
|
||
|
||
# 目前任務標記（每個 worker 流程各自一份），讓交錯的 log 可以歸屬到 (volume, level, side)
|
||
_TASK_TAG = {'vid': None, 'level': None, 'side': None}
|
||
|
||
_NP_WRAP_RE = re.compile(r'\bnp\.[A-Za-z_][A-Za-z0-9_]*\(([^()]*)\)')
|
||
|
||
|
||
def set_task_tag(volume_id, level, side=None):
|
||
_TASK_TAG['vid'] = volume_id
|
||
_TASK_TAG['level'] = level
|
||
_TASK_TAG['side'] = side
|
||
|
||
|
||
class _Tee:
|
||
"""同時輸出到 console 與 log 檔，逐行加上如 [0001 L1 LEFT] 的標記。
|
||
volume 取 id 末段（1.3.6.1.4.1.9328.50.4.0001 -> 0001）；
|
||
side 由 [LEFT]/[RIGHT]/左側/右側 段落標記判定並沿用給後續行；
|
||
最終結果段屬於整椎，重置沿用值，其中的 Left/Right 摘要行只標該行本身；
|
||
np.float64(...) 之類包裹會解包成裸數值。"""
|
||
|
||
def __init__(self, console, log_fh):
|
||
self.console = console
|
||
self.log_fh = log_fh
|
||
self.buf = ''
|
||
|
||
@staticmethod
|
||
def _marker_side(line):
|
||
"""回傳 (side, persist)：[LEFT]/[RIGHT]/左側/右側 是段落標記（persist=True），
|
||
Left/Right 摘要行只標該行（persist=False）"""
|
||
s = line.lstrip()
|
||
if s.startswith('[LEFT]') or '左側' in s:
|
||
return 'LEFT', True
|
||
if s.startswith('[RIGHT]') or '右側' in s:
|
||
return 'RIGHT', True
|
||
if s.startswith('Left '):
|
||
return 'LEFT', False
|
||
if s.startswith('Right '):
|
||
return 'RIGHT', False
|
||
return None, False
|
||
|
||
def _prefix(self, side):
|
||
if not _TASK_TAG['vid'] or not _TASK_TAG['level']:
|
||
return ''
|
||
vol = _TASK_TAG['vid'].rsplit('.', 1)[-1]
|
||
tag = f'{vol} {_TASK_TAG["level"]}'
|
||
if side:
|
||
tag += f' {side}'
|
||
return f'[{tag}] '
|
||
|
||
def _emit(self, line):
|
||
line = _NP_WRAP_RE.sub(r'\1', line)
|
||
if '最終結果' in line:
|
||
_TASK_TAG['side'] = None
|
||
side, persist = self._marker_side(line)
|
||
if persist:
|
||
_TASK_TAG['side'] = side
|
||
eff_side = side if side is not None else _TASK_TAG['side']
|
||
prefix = self._prefix(eff_side) if line.strip() else ''
|
||
if prefix:
|
||
# 前綴已含 side 時，去掉行首重複的 [LEFT] / [RIGHT] 標記
|
||
marker = f'[{eff_side}] '
|
||
if line.startswith(marker):
|
||
line = line[len(marker):]
|
||
self.console.write(prefix + line + '\n')
|
||
self.log_fh.write(prefix + 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):
|
||
"""把這個流程的 stdout/stderr 同時寫到 log_path（append、line-buffered）"""
|
||
log_fh = open(log_path, 'a', buffering=1)
|
||
sys.stdout = _Tee(sys.stdout, log_fh)
|
||
sys.stderr = _Tee(sys.stderr, log_fh)
|
||
|
||
|
||
def _setup_logging():
|
||
"""须在 setup_tee 之後呼叫，讓 handler 寫入 Tee（同時進 console 與 log 檔）"""
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S",
|
||
)
|
||
|
||
|
||
def get_device(gpu_id=None):
|
||
if torch.cuda.is_available():
|
||
if gpu_id is None:
|
||
max_free = -1
|
||
gpu_id = 0
|
||
for i in range(torch.cuda.device_count()):
|
||
free_mem, _ = torch.cuda.mem_get_info(i)
|
||
# print(f'GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)} {free_mem}')
|
||
if free_mem > max_free:
|
||
max_free = free_mem
|
||
gpu_id = i
|
||
device = torch.device(f"cuda:{gpu_id}")
|
||
logger.info(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
|
||
else:
|
||
device = torch.device("cpu")
|
||
logger.info("CUDA not available, using CPU")
|
||
return device
|
||
|
||
def debug_orientation(volume_id, level):
|
||
|
||
|
||
volume_dir = os.path.join(standardized_dir, volume_id)
|
||
|
||
# _binary.nii.gz 現為原解析度；0.5mm 用 _binary_sdf（SDF 平滑遮罩）
|
||
# 未旋轉檔：新世代在 crop/ 子資料夾、舊世代在頂層
|
||
sdf = level_file_path(volume_dir, level, 'binary_sdf')
|
||
binary_path = sdf if os.path.exists(sdf) \
|
||
else level_file_path(volume_dir, level, 'binary')
|
||
|
||
# azi = azimuth_rotation(binary_path)
|
||
# res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=False, debug=False)
|
||
azi = azimuth_rotation(binary_path, show_plt=True, save_plt=True, output_path=f'{azimuth_rotation_dir}/{level}_{volume_id}.png')
|
||
res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=True, debug=False, save_plt=True, output_path=f'{tilt_contour_dir}/{level}_{volume_id}.png')
|
||
alt = res['superior']['tilt_angle_deg']
|
||
|
||
logger.info(binary_path)
|
||
# print(f'Azimuth: {azi}, Alt: {alt}')
|
||
logger.info(f'Alt: {alt}')
|
||
|
||
def debug_pso(volume_id, level, device=None, side='both', run_id=None):
|
||
# ====== PSO ======
|
||
swarm_size = 100
|
||
max_iter = 100
|
||
|
||
# ====== DEVICE ======
|
||
if device is None:
|
||
device = get_device()
|
||
|
||
# ====== OTHER ======
|
||
spacing = [0.5, 0.5, 0.5]
|
||
CBT = True
|
||
|
||
volume_dir = os.path.join(standardized_dir, volume_id)
|
||
# _cortical 現只產出旋轉版（rotated/ 子資料夾）；_roi2 不再存檔，
|
||
# 改用旋轉後 _roi（CT）。
|
||
rotated_dir = os.path.join(volume_dir, 'rotated')
|
||
cortical_path = os.path.join(rotated_dir, f'{level}_cortical.nii.gz')
|
||
binary_path = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz')
|
||
roi_path = os.path.join(rotated_dir, f'{level}_roi.nii.gz')
|
||
|
||
missing = [p for p in (cortical_path, binary_path, roi_path) if not os.path.exists(p)]
|
||
if missing:
|
||
raise ValueError(f'{volume_id} {level}: missing {missing}')
|
||
|
||
cortical_image = sitk.ReadImage(cortical_path)
|
||
binary_image = sitk.ReadImage(binary_path)
|
||
roi_image = sitk.ReadImage(roi_path)
|
||
|
||
cortical_array = sitk.GetArrayFromImage(cortical_image)
|
||
binary_array = sitk.GetArrayFromImage(binary_image)
|
||
roi_array = sitk.GetArrayFromImage(roi_image)
|
||
image_shape = binary_array.shape
|
||
|
||
# 資料層級的快速失敗（例：…9328.50.4.0653 L5 是 1-voxel 寬的退化體積）
|
||
if binary_array.sum() == 0:
|
||
raise ValueError(f'{volume_id} {level}: empty bone mask')
|
||
if min(image_shape) < 8:
|
||
raise ValueError(f'{volume_id} {level}: degenerate volume shape {image_shape}')
|
||
|
||
cortical_tensor = torch.tensor(cortical_array, device=device)
|
||
binary_tensor = torch.tensor(binary_array, device=device)
|
||
|
||
grid = create_coordinate_grid(image_shape, device)
|
||
|
||
set_global_context(
|
||
cortical=cortical_tensor,
|
||
spine=binary_tensor,
|
||
shape=image_shape,
|
||
spacing_=spacing,
|
||
device_=device,
|
||
grid_=grid,
|
||
use_tip_penalty=False
|
||
)
|
||
|
||
best_l, loss_l, best_r, loss_r, total_time = run_pso_torch_xfr(
|
||
label_str=f'{volume_id} {level}',
|
||
image1_path=cortical_path,
|
||
image2_path=binary_path,
|
||
image3_path=roi_path,
|
||
folder=Output_dir,
|
||
swarm_size=swarm_size,
|
||
max_iter=max_iter,
|
||
spacing=spacing,
|
||
CBT=CBT,
|
||
device=device,
|
||
optimize_size=True,
|
||
grid=grid,
|
||
side=side,
|
||
level=level,
|
||
patient_id=volume_id,
|
||
side_dir=SIDE_RESULT_DIR,
|
||
run_id=run_id,
|
||
)
|
||
|
||
# exit()
|
||
|
||
|
||
def list_gpu_ids():
|
||
"""透過 nvidia-smi 取得 GPU ID（主流程不初始化 CUDA，避免污染 fork/spawn 子流程）"""
|
||
try:
|
||
out = subprocess.check_output(
|
||
['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'],
|
||
text=True,
|
||
)
|
||
return [int(line.strip()) for line in out.splitlines() if line.strip()]
|
||
except Exception:
|
||
return []
|
||
|
||
|
||
def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue):
|
||
"""每張 GPU 一個工作流程：先鎖死該 GPU，再從共享隊列領 (volume, level) 任務"""
|
||
# worker 是獨立流程，要自己把輸出 tee 到 log 檔
|
||
setup_tee(log_path)
|
||
_setup_logging()
|
||
# 必須在任何 torch.cuda 呼叫前設定
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
|
||
torch.cuda.set_device(0)
|
||
device = torch.device('cuda:0')
|
||
logger.info(f'=== [GPU {gpu_id}] worker started ===')
|
||
|
||
while True:
|
||
# 注意：multiprocessing.Queue 沒有 task_done()（只有 queue.Queue 有），別加回來
|
||
item = task_queue.get()
|
||
if item is None:
|
||
break
|
||
volume_id, level, side = item
|
||
# tag 沿用 _Tee 的 LEFT/RIGHT 寫法（'L'/'R' 是 optimizer 的 side 值）
|
||
set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT')
|
||
try:
|
||
debug_pso(volume_id, level, device, side=side, run_id=run_id)
|
||
result_queue.put(('task', gpu_id, volume_id, level, side, True, ''))
|
||
except Exception as e:
|
||
logger.error(f'[GPU {gpu_id}] Error in {volume_id} {level} {side}: {e}')
|
||
result_queue.put(('task', gpu_id, volume_id, level, side, False, str(e)))
|
||
|
||
result_queue.put(('done', gpu_id))
|
||
logger.info(f'=== [GPU {gpu_id}] worker finished ===')
|
||
|
||
|
||
def _run_sequential(tasks, run_id):
|
||
"""沒有（或只有一張）GPU 時的回退：單流程串行"""
|
||
device = get_device()
|
||
for volume_id, level, side in tasks:
|
||
set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT')
|
||
try:
|
||
debug_pso(volume_id, level, device, side=side, run_id=run_id)
|
||
except Exception as e:
|
||
logger.error(f'Error in {volume_id} {level} {side}: {e}')
|
||
|
||
|
||
USAGE = 'Usage: python xfr_debug.py [volume_id] [level]'
|
||
|
||
|
||
def parse_args(argv):
|
||
"""volume_id 可用完整 ID 或末段（如 0005）；level 為 LEVELS 之一（L1~L5）。
|
||
兩者可省略（＝全部）；level 必須搭配 volume_id 使用。"""
|
||
vid_arg = argv[0] if len(argv) >= 1 else None
|
||
level_arg = argv[1] if len(argv) >= 2 else None
|
||
if len(argv) > 2:
|
||
sys.exit(f'{USAGE}\nToo many arguments')
|
||
if level_arg and level_arg.upper() not in LEVELS:
|
||
sys.exit(f'{USAGE}\nInvalid level: {level_arg} (choose from {"/".join(LEVELS)})')
|
||
if level_arg and not vid_arg:
|
||
sys.exit(f'{USAGE}\nlevel requires volume_id')
|
||
return vid_arg, (level_arg.upper() if level_arg else None)
|
||
|
||
|
||
def main():
|
||
# level = 'L1'
|
||
|
||
# volume_id = '1.3.6.1.4.1.9328.50.4.0001'
|
||
# volume_id = '1.3.6.1.4.1.9328.50.4.0003'
|
||
# # volume_id = '1.3.6.1.4.1.9328.50.4.0121'
|
||
# process_volume(volume_id, level)
|
||
# exit()
|
||
|
||
# 要處理的 volume 數（並行模式下是「最多嘗試的 volume 數」）
|
||
MAX_SUCCESSFUL_VOLUMES = 100
|
||
MAX_SUCCESSFUL_VOLUMES = 10
|
||
|
||
# log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔）
|
||
os.makedirs(LOG_DIR, exist_ok=True)
|
||
run_id = time.strftime("%Y%m%d_%H%M%S")
|
||
log_path = os.path.join(LOG_DIR, f'xfr_debug_{run_id}.log')
|
||
setup_tee(log_path)
|
||
_setup_logging()
|
||
logger.info(f'Log file: {log_path}')
|
||
logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}')
|
||
logger.info(f'Working directory: {os.getcwd()}')
|
||
|
||
volumes = [d for d in sorted(os.listdir(standardized_dir))
|
||
if os.path.isdir(os.path.join(standardized_dir, d))]
|
||
|
||
vid_arg, level_arg = parse_args(sys.argv[1:])
|
||
if vid_arg is not None:
|
||
key = vid_arg.lower()
|
||
vols = [v for v in volumes
|
||
if v.lower() == key or v.rsplit('.', 1)[-1] == key]
|
||
if not vols:
|
||
sys.exit(f'Volume not found: {vid_arg}')
|
||
volumes = vols
|
||
volumes = volumes[:MAX_SUCCESSFUL_VOLUMES]
|
||
|
||
levels = (level_arg,) if level_arg else LEVELS
|
||
# 任務粒度 = (volume, level, side)：同一 level 的 L/R 是兩個獨立任務，
|
||
# 可被不同 GPU 的 worker 領走並行執行；順序 L1 L -> L1 R -> L2 L -> L2 R -> ...
|
||
tasks = [(vid, level, s) for vid in volumes for level in levels for s in ('L', 'R')]
|
||
logger.info(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level, side) tasks '
|
||
f'(levels: {", ".join(levels)})')
|
||
if vid_arg or level_arg:
|
||
logger.info(f'Filter: volume_id={vid_arg!r} level={level_arg!r}')
|
||
|
||
gpu_ids = list_gpu_ids()
|
||
|
||
if len(gpu_ids) <= 1:
|
||
logger.info(f'Only {len(gpu_ids)} GPU(s) available, running sequentially')
|
||
_run_sequential(tasks, run_id)
|
||
return
|
||
|
||
logger.info(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)')
|
||
|
||
ctx = mp.get_context('spawn')
|
||
task_queue = ctx.Queue()
|
||
result_queue = ctx.Queue()
|
||
for t in tasks:
|
||
task_queue.put(t)
|
||
for _ in gpu_ids:
|
||
task_queue.put(None) # 每個 worker 一個結束哨兵
|
||
|
||
procs = [ctx.Process(target=gpu_worker, args=(g, log_path, run_id, task_queue, result_queue), name=f'cbt-gpu-{g}')
|
||
for g in gpu_ids]
|
||
for p in procs:
|
||
p.start()
|
||
|
||
start_time = time.time()
|
||
results = []
|
||
finished = 0
|
||
while finished < len(gpu_ids):
|
||
if not any(p.is_alive() for p in procs):
|
||
logger.warning('A worker exited early; reaping remaining tasks...')
|
||
break
|
||
try:
|
||
msg = result_queue.get(timeout=5)
|
||
except queue_module.Empty:
|
||
continue
|
||
if msg[0] == 'done':
|
||
finished += 1
|
||
else:
|
||
results.append(msg)
|
||
|
||
# 抽乾剩下排進來的結果
|
||
while True:
|
||
try:
|
||
msg = result_queue.get(timeout=1)
|
||
except queue_module.Empty:
|
||
break
|
||
if msg[0] == 'task':
|
||
results.append(msg)
|
||
|
||
for p in procs:
|
||
p.join(timeout=60)
|
||
|
||
total_time = time.time() - start_time
|
||
ok = [r for r in results if r[5]]
|
||
fail = [r for r in results if not r[5]]
|
||
per_volume = {}
|
||
per_volume_n = {}
|
||
for _, _, vid, level, side, success, _ in results:
|
||
per_volume.setdefault(vid, set()).add(success)
|
||
per_volume_n[vid] = per_volume_n.get(vid, 0) + 1
|
||
missing = len(tasks) - len(results)
|
||
|
||
# 一個 volume 算「成功」必須它的所有 (level, side) 任務都執行過且全部成功
|
||
n_success_volumes = sum(
|
||
1 for vid in volumes
|
||
if per_volume.get(vid, set()) == {True}
|
||
and per_volume_n.get(vid, 0) == len(levels) * 2)
|
||
|
||
print('=' * 60)
|
||
logger.info(f'Finished in {total_time / 60:.1f} min | '
|
||
f'tasks {len(results)}/{len(tasks)} (ok {len(ok)} / failed {len(fail)} / not-run {missing})')
|
||
if missing:
|
||
logger.warning(f'{missing} task(s) were never executed (worker crash?)')
|
||
logger.info(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}')
|
||
if fail:
|
||
logger.info('Failed tasks:')
|
||
for _, g, vid, level, side, _, err in fail:
|
||
logger.error(f'[GPU {g}] {vid} {level} {side}: {err}')
|
||
|
||
# 收尾：螺絲位置映回原 CT 空間 -> Output_dir/<run_date>/<volume_id>/cbt.nii.gz
|
||
# （label 1-10 = L1L L1R L2L L2R ... L5L L5R；無 side 結果的 volume 跳過）
|
||
for vid in volumes:
|
||
try:
|
||
xfr_cbt_native.write_volume_cbt(vid, run_id)
|
||
except Exception as e:
|
||
logger.error(f'[CBT-NATIVE] {vid}: {e}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |