Xiao Furen
4f5be3d3e9
Refactor the optimization pipeline to eliminate module-level global variables, improving thread safety and modularity. Introduced `OptimizationContext` to explicitly manage shared state during cylinder evaluation. Key changes: - core: Replace global variables with `OptimizationContext` dataclass in `objective.py`. - core: Implement `refine_lateral_longer` in `optimizer.py` for deterministic local refinement of screw placement. - core: Update scoring logic in `scoring.py` to use higher penalties for out-of-bone voxels. - imaging: Add advanced symmetry detection including `best_symmetry_plane` and `best_symmetry_axis_angle` in `orientation.py`. - visualization: Enhance 3D plotting in `res_plot_3d.py` with volume absorption rendering (Beer-Lambert law) for an X-ray-like appearance. - xfr_debug: Implement a custom `_Tee` logger to support multi-process logging with volume and level-specific tags. - chore: Update `.gitignore` to include local logs and kilo directories.
384 lines
No EOL
13 KiB
Python
384 lines
No EOL
13 KiB
Python
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
|
||
|
||
standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/'
|
||
|
||
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')
|
||
|
||
|
||
# 目前任務標記（每個 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):
|
||
_TASK_TAG['vid'] = volume_id
|
||
_TASK_TAG['level'] = level
|
||
_TASK_TAG['side'] = None
|
||
|
||
|
||
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 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}")
|
||
print(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
|
||
else:
|
||
device = torch.device("cpu")
|
||
print("CUDA not available, using CPU")
|
||
return device
|
||
|
||
def debug_orientation(volume_id, level):
|
||
|
||
|
||
volume_dir = os.path.join(standardized_dir, volume_id)
|
||
|
||
cortical_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz')
|
||
binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz')
|
||
roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz')
|
||
|
||
# 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']
|
||
|
||
print(binary_path)
|
||
# print(f'Azimuth: {azi}, Alt: {alt}')
|
||
print(f'Alt: {alt}')
|
||
|
||
def debug_pso(volume_id, level, device=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_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz')
|
||
binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz')
|
||
roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz')
|
||
|
||
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,
|
||
)
|
||
|
||
# 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, task_queue, result_queue):
|
||
"""每張 GPU 一個工作流程：先鎖死該 GPU，再從共享隊列領 (volume, level) 任務"""
|
||
# worker 是獨立流程，要自己把輸出 tee 到 log 檔
|
||
setup_tee(log_path)
|
||
# 必須在任何 torch.cuda 呼叫前設定
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
|
||
torch.cuda.set_device(0)
|
||
device = torch.device('cuda:0')
|
||
print(f'=== [GPU {gpu_id}] worker started ===', flush=True)
|
||
|
||
while True:
|
||
# 注意：multiprocessing.Queue 沒有 task_done()（只有 queue.Queue 有），別加回來
|
||
item = task_queue.get()
|
||
if item is None:
|
||
break
|
||
volume_id, level = item
|
||
set_task_tag(volume_id, level)
|
||
try:
|
||
debug_pso(volume_id, level, device)
|
||
result_queue.put(('task', gpu_id, volume_id, level, True, ''))
|
||
except Exception as e:
|
||
print(f'[GPU {gpu_id}] Error in {volume_id} {level}: {e}', flush=True)
|
||
result_queue.put(('task', gpu_id, volume_id, level, False, str(e)))
|
||
|
||
result_queue.put(('done', gpu_id))
|
||
print(f'=== [GPU {gpu_id}] worker finished ===', flush=True)
|
||
|
||
|
||
def _run_sequential(tasks):
|
||
"""沒有（或只有一張）GPU 時的回退：單流程串行"""
|
||
device = get_device()
|
||
for volume_id, level in tasks:
|
||
set_task_tag(volume_id, level)
|
||
try:
|
||
debug_pso(volume_id, level, device)
|
||
except Exception as e:
|
||
print(f'Error in {volume_id} {level}: {e}')
|
||
|
||
|
||
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 = 1
|
||
|
||
# log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔）
|
||
os.makedirs(LOG_DIR, exist_ok=True)
|
||
log_path = os.path.join(LOG_DIR, f'xfr_debug_{time.strftime("%Y%m%d_%H%M%S")}.log')
|
||
setup_tee(log_path)
|
||
print(f'Log file: {log_path}', flush=True)
|
||
print(f'Command: {sys.executable} {" ".join(sys.argv)}', flush=True)
|
||
print(f'Working directory: {os.getcwd()}', flush=True)
|
||
|
||
volumes = [d for d in sorted(os.listdir(standardized_dir))
|
||
if os.path.isdir(os.path.join(standardized_dir, d))]
|
||
volumes = volumes[:MAX_SUCCESSFUL_VOLUMES]
|
||
|
||
tasks = [(vid, level) for vid in volumes for level in LEVELS]
|
||
print(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level) tasks', flush=True)
|
||
|
||
gpu_ids = list_gpu_ids()
|
||
|
||
if len(gpu_ids) <= 1:
|
||
print(f'Only {len(gpu_ids)} GPU(s) available, running sequentially', flush=True)
|
||
_run_sequential(tasks)
|
||
return
|
||
|
||
print(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)', flush=True)
|
||
|
||
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, 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):
|
||
print('Warning: a worker exited early; reaping remaining tasks...', flush=True)
|
||
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[4]]
|
||
fail = [r for r in results if not r[4]]
|
||
per_volume = {}
|
||
for _, _, vid, level, success, _ in results:
|
||
per_volume.setdefault(vid, set()).add(success)
|
||
missing = len(tasks) - len(results)
|
||
|
||
# 一個 volume 算「成功」必須它的(level)全部執行過且全部成功
|
||
n_success_volumes = sum(1 for vid in volumes
|
||
if per_volume.get(vid, set()) == {True})
|
||
|
||
print('=' * 60)
|
||
print(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:
|
||
print(f'Warning: {missing} task(s) were never executed (worker crash?)')
|
||
print(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}')
|
||
if fail:
|
||
print('Failed tasks:')
|
||
for _, g, vid, level, _, err in fail:
|
||
print(f' [GPU {g}] {vid} {level}: {err}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |