Xiao Furen
4204d2cd4c
Implement a more robust scoring mechanism for screw optimization and add functionality to generate synthetic X-ray projections (AP and lateral views) from CT data. Key changes: - core: add `generate_cylinder_butt_torch` to create a mask for the screw entrance (0.25mm) to exempt it from bone-breaching penalties. - core: update `cl_score_torch_xfr` to include a diameter preference bonus and utilize the entrance mask. - core: adjust optimizer bounds and scoring weights to favor larger diameter screws and improve convergence. - xfr_cbt_native: implement `render_xray_projections` to generate synthetic AP and lateral X-ray images for visualization. - visualization: enhance `render_bone_figure` with semi-transparent spinous process rendering and improved depth sorting for screws. - xfr_debug: improve level detection to support arbitrary lumbar levels (L1-L9) and add safe volume-level cleanup for CBT writing. - config: update allowed diameters and lengths constants.
545 lines
No EOL
20 KiB
Python
545 lines
No EOL
20 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
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
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'
|
||
|
||
def available_levels(volume_id):
|
||
"""該 volume 可跑 debug_pso 的全部 lumbar level：rotated/ 中輸入三件
|
||
（_cortical / _binary_sdf / _roi）齊全的 L\\d。各 volume 層數不同
|
||
（例：有的只有 L1~L3、有的含 L6），故不用固定 LEVELS 清單。"""
|
||
rotated_dir = os.path.join(standardized_dir, volume_id, 'rotated')
|
||
if not os.path.isdir(rotated_dir):
|
||
return ()
|
||
levels = []
|
||
for fn in os.listdir(rotated_dir):
|
||
m = re.fullmatch(r'(L[1-9]\d*)_cortical\.nii\.gz', fn)
|
||
if not m:
|
||
continue
|
||
level = m.group(1)
|
||
if (os.path.exists(os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz'))
|
||
and os.path.exists(os.path.join(rotated_dir, f'{level}_roi.nii.gz'))):
|
||
levels.append(level)
|
||
return tuple(sorted(levels, key=lambda lv: int(lv[1:])))
|
||
|
||
|
||
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 _write_cbt_safe(volume_id, run_id):
|
||
"""per-volume 收尾：cbt.nii.gz + x-ap/x-lat 投影（錯誤已 log，回傳 ok）"""
|
||
try:
|
||
xfr_cbt_native.write_volume_cbt(volume_id, run_id)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f'[CBT-NATIVE] {volume_id}: {e}')
|
||
return False
|
||
|
||
|
||
def _run_sequential(tasks, run_id):
|
||
"""沒有（或只有一張）GPU 時的回退：單流程串行。
|
||
tasks 依 (volume, level, side) 分組排列：一個 volume 的 (level, side)
|
||
全部跑完後，立刻寫它的 cbt.nii.gz + 投影（不等其他 volume）。"""
|
||
device = get_device()
|
||
current_vid = None
|
||
for volume_id, level, side in tasks:
|
||
if volume_id != current_vid:
|
||
if current_vid is not None:
|
||
_write_cbt_safe(current_vid, run_id)
|
||
current_vid = volume_id
|
||
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}')
|
||
if current_vid is not None:
|
||
_write_cbt_safe(current_vid, run_id)
|
||
|
||
|
||
USAGE = 'Usage: python xfr_debug.py [volume_id] [level]'
|
||
|
||
|
||
def parse_args(argv):
|
||
"""volume_id 可用完整 ID 或末段（如 0005）；level 為 L\\d 形式（L1、L2、…，
|
||
不限定 L1~L5；實際執行以各 volume 資料中有的 level 為準，見 available_levels）。
|
||
兩者可省略（＝全部）；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 not re.fullmatch(r'L[1-9]\d*', level_arg.upper()):
|
||
sys.exit(f'{USAGE}\nInvalid level: {level_arg} (expected L1, L2, ...)')
|
||
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]
|
||
|
||
# 各 volume 的 lumbar level：未指定 level 時，取該 volume 實際有的全部 lumbar level
|
||
# （見 available_levels）；指定 level 時只跑該 level（該 volume 缺檔會直接報錯）
|
||
vol_levels = {}
|
||
for vid in volumes:
|
||
vol_levels[vid] = (level_arg,) if level_arg else available_levels(vid)
|
||
no_levels = [vid for vid in volumes if not level_arg and not vol_levels[vid]]
|
||
if no_levels:
|
||
logger.warning(f'{len(no_levels)} volume(s) have no available lumbar level, '
|
||
f'skipped: {", ".join(no_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 vol_levels[vid] for s in ('L', 'R')]
|
||
expected = {vid: len(vol_levels[vid]) * 2 for vid in volumes}
|
||
all_levels = sorted({level for lv in vol_levels.values() for level in lv},
|
||
key=lambda lv: int(lv[1:]))
|
||
logger.info(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level, side) tasks '
|
||
f'(levels: {", ".join(all_levels) if all_levels else "(none)"})')
|
||
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 一個結束哨兵
|
||
|
||
# per-volume 收尾：一個 volume 排入的 (level, side) 任務全部回報後
|
||
# （成功或失敗），立刻在背景執行緒寫它的 cbt.nii.gz + 投影，不等全部 case
|
||
write_pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix='cbt-write')
|
||
write_futures = set()
|
||
pending_writes = set(volumes)
|
||
remaining = {vid: expected[vid] for vid in volumes}
|
||
|
||
def _fire_write(vid):
|
||
pending_writes.discard(vid)
|
||
write_futures.add(write_pool.submit(_write_cbt_safe, vid, run_id))
|
||
|
||
def _on_task_result(msg):
|
||
results.append(msg)
|
||
vid = msg[2]
|
||
if vid in remaining:
|
||
remaining[vid] -= 1
|
||
if remaining[vid] <= 0 and vid in pending_writes:
|
||
logger.info(f'{vid}: screw tasks complete, writing cbt + projections now')
|
||
_fire_write(vid)
|
||
|
||
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:
|
||
_on_task_result(msg)
|
||
|
||
# 抽乾剩下排進來的結果
|
||
while True:
|
||
try:
|
||
msg = result_queue.get(timeout=1)
|
||
except queue_module.Empty:
|
||
break
|
||
if msg[0] == 'task':
|
||
_on_task_result(msg)
|
||
|
||
for p in procs:
|
||
p.join(timeout=60)
|
||
|
||
# 補漏：worker 提早退出、有任務未回報的 volume 仍照舊嘗試寫
|
||
# （無 side 結果時 write_volume_cbt 會自行 skip）
|
||
for vid in volumes:
|
||
if vid in pending_writes:
|
||
_fire_write(vid)
|
||
|
||
# 等待所有 per-volume cbt.nii.gz / 投影寫出完成
|
||
write_ok = sum(1 for f in write_futures if f.result())
|
||
write_fail = len(write_futures) - write_ok
|
||
write_pool.shutdown(wait=True)
|
||
|
||
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) 任務都執行過且全部成功
|
||
# （expected 為該 volume 實際排入的任務數，各 volume 的 level 數可不同）
|
||
n_success_volumes = sum(
|
||
1 for vid in volumes
|
||
if per_volume.get(vid, set()) == {True}
|
||
and per_volume_n.get(vid, 0) == expected.get(vid, 0))
|
||
|
||
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}')
|
||
|
||
# cbt.nii.gz + x-ap.jpg / x-lat.jpg 已於各 volume 的螺絲任務完成後立刻
|
||
# 寫出（_fire_write， Output_dir/<run_date>/<volume_id>/；label 1-10 =
|
||
# L1L L1R L2L L2R ... L5L L5R），不再等全部 case 跑完才統一收尾
|
||
logger.info(f'CBT writes: {write_ok}/{len(volumes)} volume(s) ok'
|
||
+ (f', {write_fail} failed (見 [CBT-NATIVE] log)' if write_fail else ''))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |