Xiao Furen
523ec7ee16
Introduces a new scoring component for vertebral body (VBODY) rewards to improve optimization accuracy. The update includes: - Added `vbody_tensor` to `OptimizationContext` and scoring functions to reward screw placement within the vertebral body. - Enhanced `segment_spinous_process` with diagnostic capabilities to detect spinous process absence (e.g., post-laminectomy). - Improved `resample_img` to prevent physical boundary clipping and handle interpolation more robustly for CT and label data. - Implemented a metadata cache using TinyDB in the preprocessing pipeline to skip low-resolution or insufficient scans efficiently. - Added robust error handling for NFS-based file operations and directory creation. - Added new visualization tools for bone figures and level plotting. refactor(imaging): improve segmentation and resampling precision - Refactored `seg_bone` to support original resolution binary masks and Signed Maurer Distance Maps (SMD) for more accurate boundary handling. - Updated `resample_img` to use `ceil` for output size calculation to ensure full physical coverage. - Optimized `process_single_image` to utilize metadata for skipping processing of invalid or low-quality scans.
441 lines
No EOL
20 KiB
Python
441 lines
No EOL
20 KiB
Python
#!/home/xfr/.conda/envs/cbt/bin/python
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from tinydb import TinyDB, Query
|
||
|
||
from imaging.preprocessing import process_dataset
|
||
from config.constant import LABEL_MAP
|
||
from imaging.orientation import (best_symmetry_plane, best_upper_endplate_plane,
|
||
segment_spinous_process, segment_vertebral_body,
|
||
smooth_mask_sdf)
|
||
from visualization.res_bone_figure import (render_bone_figure,
|
||
compute_normalizing_rotation,
|
||
rotate_volume_to, rotated_grid,
|
||
rotated_sitk_image_at,
|
||
_rotate_plane_params,
|
||
_shift_plane_params)
|
||
|
||
data_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
|
||
label_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/'
|
||
output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2/'
|
||
output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'
|
||
|
||
label_map = {
|
||
'colon': 'conlon',
|
||
'COVID-19': 'COVID-19',
|
||
'HNSCC-3DCT-RT_neck': 'HNSCC-3DCT-RT_neck',
|
||
'liver': 'Liver',
|
||
}
|
||
|
||
LUMBAR_LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6')
|
||
MAX_Z_SPACING_MM = 4.0 # z spacing (mm) 大於此值者跳過整支 pipeline
|
||
MIN_LUMBAR_LEVELS = 2 # 符合 lumbar 的 label 少於此值者跳過整支 pipeline
|
||
|
||
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_preprocess')
|
||
|
||
|
||
def _upsert_by_name(table, name, meta):
|
||
"""以 'name' 欄位為鍵 upsert（TinyDB 4.9 insert 不支援自訂 doc_id）。"""
|
||
q = Query().name == name
|
||
entry = dict(table.get(q) or {})
|
||
entry.update(meta)
|
||
entry['name'] = name
|
||
if table.get(q) is None:
|
||
table.insert(entry)
|
||
else:
|
||
table.update(entry, q)
|
||
|
||
|
||
def _migrate_legacy_metadata(path):
|
||
"""舊版純 JSON 檔（{name: {...}}、無 TinyDB 的 _default 結構）：
|
||
改名成 .legacy-<timestamp> 備份，內容匯入新 TinyDB。"""
|
||
if not os.path.exists(path):
|
||
return
|
||
try:
|
||
with open(path) as f:
|
||
data = json.load(f)
|
||
except (json.JSONDecodeError, OSError):
|
||
return
|
||
legacy = {k: v for k, v in data.items()
|
||
if isinstance(v, dict)} if isinstance(data, dict) else {}
|
||
if not (isinstance(data, dict) and '_default' not in data and legacy):
|
||
return
|
||
bak = f'{path}.legacy-{time.strftime("%Y%m%d_%H%M%S")}'
|
||
os.replace(path, bak)
|
||
print(f'Metadata db {path} is old pure-JSON format; renamed to {bak} '
|
||
f'and migrating {len(legacy)} entrie(s)')
|
||
t = TinyDB(path).table('images')
|
||
for k, v in legacy.items():
|
||
_upsert_by_name(t, k, dict(v))
|
||
|
||
|
||
class ImageMetadataDB:
|
||
"""影像 metadata 快取（TinyDB，table 'images'）：
|
||
以影像檔名前綴（= 輸出資料夾名，如 1.3.6.1.4.1.9328.50.4.0001、
|
||
liver_100、volume-covid19-A-0011_ct）為 key（存於文件 'name' 欄位），
|
||
記錄 {'spacing': [x, y, z], 'labels': [label id]}。
|
||
|
||
供 process_single_image 做跳過判定（z spacing / lumbar 層數），
|
||
命中時不必讀影像 / label 檔。TinyDB 每次操作自動落檔，
|
||
中途重跑可繼承。
|
||
"""
|
||
|
||
def __init__(self, path=METADATA_DB):
|
||
self.path = path
|
||
_migrate_legacy_metadata(path)
|
||
self._table = TinyDB(path).table('images')
|
||
|
||
def get(self, name):
|
||
return self._table.get(Query().name == name)
|
||
|
||
def put(self, name, meta):
|
||
_upsert_by_name(self._table, name, meta)
|
||
|
||
def __len__(self):
|
||
return len(self._table)
|
||
|
||
|
||
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):
|
||
"""把這個流程的 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 _cortical_from_roi(roi_arr, bin_arr):
|
||
"""皮質遮罩（與舊未旋轉 _cortical.nii.gz 同定義）：
|
||
門檻 = 骨頭 mask（bin_arr）內 CT 的 median HU；
|
||
cortical = mask 內且 門檻 <= HU <= 10000。
|
||
roi_arr / bin_arr 須同 grid（(z, y, x)）；回傳 uint8 0/1，
|
||
輸入缺失或 mask 為空時回傳 None。"""
|
||
if roi_arr is None or bin_arr is None:
|
||
return None
|
||
inside = bin_arr > 0
|
||
if int(inside.sum()) == 0:
|
||
return None
|
||
threshold = float(np.median(roi_arr[inside]))
|
||
return (inside & (roi_arr >= threshold) & (roi_arr <= 10000)).astype(np.uint8)
|
||
|
||
|
||
def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path):
|
||
"""對單一 level：計算對齊旋轉（鏡稱面 -> +X、上終板 normal y=0），
|
||
存到 <vol_dir>/rotated/（先旋轉 _binary_nn，再以它的非零 bbox（±4
|
||
體素）作為所有旋轉輸出的裁切 grid，尺寸不求同未旋轉版、整顆骨頭
|
||
保留，origin 隨裁切平移）：
|
||
- _smd_resampled.nii.gz 三线性旋轉後 SMD（內負/外正）
|
||
- _binary_sdf.nii.gz 旋轉 SMD 於 0.5 閾值（與未旋轉版一致）
|
||
- _binary_nn.nii.gz NN 旋轉後 0/1 遮罩（對比用，裁切基準）
|
||
- _roi.nii.gz
|
||
- _cortical.nii.gz 旋轉 CT 以骨頭 mask 內 median HU 為門檻
|
||
（取代舊的未旋轉 _cortical，定義相同）
|
||
再畫「rotated」平面圖，並在旋轉體上做 VBODY / 棘突分割存 label map
|
||
（1=VBODY、2=棘突、3=other bone、0=background）。
|
||
平面參數經 R 剛性旋轉並換元到輸出 grid 的局部座標。
|
||
|
||
旋轉來源 + 幾何模板：patient 未旋轉的 0.5mm _smd_resampled.nii.gz
|
||
（float、裁物件 bbox，SMD < 0.5 = 內部，與 _binary_sdf 同幾何）；
|
||
缺失或格式不符時回退 _binary_nn.nii.gz（0.5mm 最近邻 0/1），
|
||
此情形只產出 _binary_nn（+ roi / cortical / 圖 / label）。
|
||
mask_path（未旋轉 0.5mm 0/1 遮罩，如 _binary_sdf）只供畫圖
|
||
（render_bone_figure 以 uint8 讀取；float 檔會被讀成全 0），
|
||
缺失時自動回退 _binary_nn。
|
||
"""
|
||
volume_id = os.path.basename(vol_dir)
|
||
|
||
template = None
|
||
smd_arr = None
|
||
if smd_path is not None and 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:
|
||
template = sitk.ReadImage(smd_path) # 0.5mm 幾何模板
|
||
smd_arr = sa
|
||
else:
|
||
logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled 不可用 '
|
||
f'(value)，改用 _binary_nn')
|
||
if template is None:
|
||
nn_in_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz')
|
||
if not os.path.exists(nn_in_path):
|
||
logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled / _binary_nn 皆不可用, skip')
|
||
return
|
||
template = sitk.ReadImage(nn_in_path)
|
||
|
||
# 平面估計輸入：SMD 時取 0.5 閾值內部，與 _binary_sdf 同定義
|
||
if smd_arr is not None:
|
||
bin_arr = (smd_arr < 0.5).astype(np.uint8)
|
||
else:
|
||
bin_arr = (sitk.GetArrayFromImage(template) > 0).astype(np.uint8)
|
||
if int(bin_arr.sum()) == 0:
|
||
logger.warning(f'[rotated] {volume_id} {level}: empty binary, skip')
|
||
return
|
||
|
||
# 圖用遮罩須與模板同幾何（SDF 與 SMD 皆裁 bbox2；NN fallback 只有
|
||
# NN 遮罩一致）；0.5mm 0/1，float 檔會被 uint8 讀取成全 0
|
||
nn_in_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz')
|
||
if smd_arr is None:
|
||
mask_path = nn_in_path if os.path.exists(nn_in_path) else mask_path
|
||
if mask_path is None or not os.path.exists(mask_path):
|
||
mask_path = nn_in_path if os.path.exists(nn_in_path) else None
|
||
|
||
sym = best_symmetry_plane(bin_arr)
|
||
symp = best_upper_endplate_plane(bin_arr)
|
||
R, c_xyz = compute_normalizing_rotation(bin_arr, sym, symp)
|
||
|
||
# 輸出 grid：未旋轉緊密 bbox 的 8 角點旋轉後的最小包圍盒 + margin——
|
||
# 尺寸不受未旋轉版約束，保證旋轉後整顆骨頭保留（同尺寸旋轉切角落）
|
||
start, size = rotated_grid(bin_arr.shape, R, c_xyz, margin=4)
|
||
|
||
rotated_dir = os.path.join(vol_dir, 'rotated')
|
||
os.makedirs(rotated_dir, exist_ok=True)
|
||
|
||
# 先旋轉 _binary_nn（order=0，未旋轉 _binary_nn 先重採樣到模板 grid
|
||
# 再旋轉），以其非零 bbox（外扩 4 體素）定所有旋轉輸出的裁切 grid
|
||
# （fstart/fsize，模板 index 系）；NN 缺失 / 為空時回退緊密包圍盒
|
||
rot_nn = None
|
||
fstart = (int(start[0]), int(start[1]), int(start[2]))
|
||
fsize = (int(size[0]), int(size[1]), int(size[2]))
|
||
if os.path.exists(nn_in_path):
|
||
nn_img = sitk.ReadImage(nn_in_path)
|
||
nn_arr = (sitk.GetArrayFromImage(
|
||
sitk.Resample(nn_img, template, interpolator=sitk.sitkNearestNeighbor,
|
||
defaultPixelValue=0)) > 0).astype(np.uint8)
|
||
rot_nn_full = (rotate_volume_to(nn_arr, R, c_xyz, start, size,
|
||
order=0, cval=0.0) > 0.5).astype(np.uint8)
|
||
p_nn = os.path.join(rotated_dir, f'{level}_binary_nn.nii.gz')
|
||
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)
|
||
rot_nn = rot_nn_full[z0:z1 + 1, y0:y1 + 1, x0:x1 + 1]
|
||
else:
|
||
rot_nn = rot_nn_full
|
||
sitk.WriteImage(rotated_sitk_image_at(template, rot_nn, fstart), p_nn)
|
||
logger.info(f'[rotated] saved {p_nn} (grid={list(fsize)}, start={fstart})')
|
||
|
||
# 旋轉後 SMD：未旋轉 _smd_resampled 的 SMD 場經三线性旋轉得以保留
|
||
# （float32 來源，map_coordinates 輸出維持浮點，<0.5 即為真閾值）；
|
||
# 體外填充值取正 SMD（背景側），避免裁切邊界出現假的閾值穿越
|
||
rot_sdf_bin = None
|
||
if smd_arr is not None:
|
||
rot_smd = rotate_volume_to(smd_arr, R, c_xyz, fstart, fsize,
|
||
order=1, cval=float(smd_arr.max()))
|
||
p_smd = os.path.join(rotated_dir, f'{level}_smd_resampled.nii.gz')
|
||
sitk.WriteImage(rotated_sitk_image_at(template, rot_smd.astype(np.float32),
|
||
fstart), p_smd)
|
||
logger.info(f'[rotated] saved {p_smd}')
|
||
|
||
# _binary_sdf.nii.gz：旋轉 SMD 於 0.5 閾值（與未旋轉版一致）
|
||
rot_sdf_bin = (rot_smd < 0.5).astype(np.uint8)
|
||
p_sdf = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz')
|
||
sitk.WriteImage(rotated_sitk_image_at(template, rot_sdf_bin, fstart), p_sdf)
|
||
logger.info(f'[rotated] saved {p_sdf}')
|
||
|
||
# 旋轉後的 roi（三线性）：先重採樣到模板 grid，再旋轉到裁切 grid
|
||
roi_arr = None
|
||
rot_roi_arr = None
|
||
if roi_path is not None and os.path.exists(roi_path):
|
||
roi_img = sitk.ReadImage(roi_path)
|
||
roi_arr = sitk.GetArrayFromImage(
|
||
sitk.Resample(roi_img, template, defaultPixelValue=0.0))
|
||
rot_roi = rotate_volume_to(roi_arr, R, c_xyz, fstart, fsize,
|
||
order=1, cval=0.0)
|
||
p_roi = os.path.join(rotated_dir, f'{level}_roi.nii.gz')
|
||
sitk.WriteImage(rotated_sitk_image_at(template, rot_roi, fstart), p_roi)
|
||
logger.info(f'[rotated] saved {p_roi}')
|
||
rot_roi_arr = rot_roi
|
||
|
||
# 分割 / cortical 輸入：SMD 時直接用旋轉 _binary_sdf（0.5 閾值、
|
||
# 次體素平滑）；NN fallback 時以 smooth_mask_sdf 記憶體內平滑（不存檔）
|
||
if rot_sdf_bin is not None:
|
||
rot_bin_seg = rot_sdf_bin
|
||
else:
|
||
rot_bin_seg = smooth_mask_sdf(rot_nn, sigma=1.5).astype(np.uint8)
|
||
|
||
# 旋轉後的 cortical：旋轉 CT 以骨頭 mask 內 median HU 為門檻
|
||
#（取代舊的未旋轉 _cortical.nii.gz，定義相同）
|
||
rot_cort = _cortical_from_roi(rot_roi_arr, rot_bin_seg)
|
||
if rot_cort is not None:
|
||
p_cort = os.path.join(rotated_dir, f'{level}_cortical.nii.gz')
|
||
sitk.WriteImage(rotated_sitk_image_at(template, rot_cort, fstart), p_cort)
|
||
logger.info(f'[rotated] saved {p_cort}')
|
||
else:
|
||
logger.warning(f'[rotated] {volume_id} {level}: 無旋轉 CT / mask，跳過 _cortical')
|
||
|
||
# 用旋轉後的平面畫圖（rotated 版 planes）；皮質著色由未旋轉 CT + mask
|
||
# 現算（未旋轉 _cortical 不再存檔）
|
||
p_fig = os.path.join(rotated_dir, f'{level}_planes.png')
|
||
if mask_path is not None:
|
||
fig_cortical = _cortical_from_roi(roi_arr, bin_arr)
|
||
fig = render_bone_figure(volume_id, level, mask_path, fig_cortical,
|
||
planes_only=True, rotation=(R, c_xyz), output_path=p_fig)
|
||
if fig is not None:
|
||
logger.info(f'[rotated] saved {fig}')
|
||
else:
|
||
logger.warning(f'[rotated] {volume_id} {level}: 無可用 0/1 遮罩，跳過平面圖')
|
||
|
||
# 旋轉體上的 VBODY / 棘突分割：平面參數隨 R 剛性旋轉（rotated 體與
|
||
# 原始體只差一個剛性變換，平面隨動即可，不需重新偵測），
|
||
# label map 存同資料夾：1 = VBODY（椎體），2 = 棘突（spinous process），
|
||
# 3 = other bone（其餘骨頭），0 = background。
|
||
# 換元到輸出 grid 的局部座標（origin 在模板 index fstart 處）
|
||
sym_rot = _shift_plane_params(_rotate_plane_params(sym, R, c_xyz), fstart)
|
||
symp_rot = (_shift_plane_params(_rotate_plane_params(symp, R, c_xyz), fstart)
|
||
if symp is not None else None)
|
||
sp_mask, sp_th, sp_info = segment_spinous_process(rot_bin_seg, sym_rot)
|
||
vb_mask, vb_th, vb_info = segment_vertebral_body(rot_bin_seg, sym_rot, symp_rot,
|
||
sp_th, sp_info['mode'])
|
||
label_arr = np.zeros(rot_bin_seg.shape, dtype=np.uint8)
|
||
label_arr[rot_bin_seg > 0] = 3 # other bone
|
||
if vb_mask is not None:
|
||
label_arr[vb_mask] = 1 # VBODY
|
||
if sp_mask is not None:
|
||
label_arr[sp_mask] = 2 # spinous process
|
||
p_lbl = os.path.join(rotated_dir, f'{level}_label.nii.gz')
|
||
sitk.WriteImage(rotated_sitk_image_at(template, label_arr, fstart), p_lbl)
|
||
logger.info(f'[rotated] saved {p_lbl} '
|
||
f'(vbody={int((label_arr == 1).sum())}, spinous={int((label_arr == 2).sum())}, '
|
||
f'other={int((label_arr == 3).sum())}, '
|
||
f'sp_mode={sp_info["mode"]}, vb_mode={vb_info["mode"]})')
|
||
|
||
|
||
def make_lumbar_post_process():
|
||
"""每個 volume 處理完後，對其 lumbar level：
|
||
1) 畫「骨頭 + 方向平面」圖（不畫螺絲、不做棘突 / 椎體分割）-> <volume_dir>/lumbar/
|
||
2) 計算對齊旋轉，存旋轉後的 smd_resampled / binary_sdf / binary_nn / roi
|
||
+ cortical + 旋轉平面圖 + label map -> <volume_dir>/rotated/
|
||
|
||
post_process 由 process_dataset 呼叫：(volume_dir, processed_labels)。
|
||
processed_labels 為該 volume 實際存在的 label id（int），對照 LABEL_MAP。
|
||
"""
|
||
|
||
def _post_process(vol_dir, processed_labels):
|
||
volume_id = os.path.basename(vol_dir)
|
||
lumbar_dir = os.path.join(vol_dir, 'lumbar')
|
||
for n in processed_labels:
|
||
level = LABEL_MAP.get(int(n))
|
||
if level not in LUMBAR_LEVELS:
|
||
continue
|
||
smd_res_path = os.path.join(vol_dir, f'{level}_smd_resampled.nii.gz')
|
||
nn_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz')
|
||
if not os.path.exists(smd_res_path) and not os.path.exists(nn_path):
|
||
continue
|
||
# 畫圖用 0/1 遮罩：優 _binary_sdf（0.5mm 平滑），缺則 _binary_nn
|
||
sdf_path = os.path.join(vol_dir, f'{level}_binary_sdf.nii.gz')
|
||
mask_path = sdf_path if os.path.exists(sdf_path) else nn_path
|
||
roi_path = os.path.join(vol_dir, f'{level}_roi.nii.gz')
|
||
|
||
# 1) 原始（未旋轉）planes 圖；皮質著色由未旋轉 CT + mask 現算
|
||
#（未旋轉 _cortical 不再存檔）
|
||
output_path = os.path.join(lumbar_dir, f'{level}_planes.png')
|
||
if os.path.exists(mask_path):
|
||
fig_cortical = None
|
||
if os.path.exists(roi_path):
|
||
mask_img = sitk.ReadImage(mask_path)
|
||
roi_img = sitk.ReadImage(roi_path)
|
||
roi_arr = sitk.GetArrayFromImage(
|
||
sitk.Resample(roi_img, mask_img, defaultPixelValue=0.0))
|
||
fig_cortical = _cortical_from_roi(
|
||
roi_arr, sitk.GetArrayFromImage(mask_img))
|
||
path = render_bone_figure(volume_id, level, mask_path, fig_cortical,
|
||
planes_only=True, output_path=output_path)
|
||
if path is not None:
|
||
logger.info(f'[lumbar] saved {path}')
|
||
|
||
# 2) 旋轉對齊：rotated/ 的 smd_resampled + binary_sdf + binary_nn
|
||
# + roi + cortical + planes 圖 + label
|
||
_write_rotated_level(vol_dir, level, smd_res_path, mask_path, roi_path)
|
||
|
||
return _post_process
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='Preprocess CT spine dataset.')
|
||
parser.add_argument('--max-images', type=int, default=None, dest='max_images',
|
||
help='Process at most this number of images per dataset (default: all).')
|
||
args = parser.parse_args()
|
||
|
||
# log 檔（console 與檔案同時輸出）
|
||
os.makedirs(LOG_DIR, exist_ok=True)
|
||
log_path = os.path.join(LOG_DIR, f'xfr_preprocess_{time.strftime("%Y%m%d_%H%M%S")}.log')
|
||
setup_tee(log_path)
|
||
# basicConfig 在 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",
|
||
)
|
||
logger.info(f'Log file: {log_path}')
|
||
logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}')
|
||
logger.info(f'Working directory: {os.getcwd()}')
|
||
|
||
# metadata db：跳過判定（z spacing / lumbar 層數）命中時免讀影像 / label 檔
|
||
metadata_db = ImageMetadataDB()
|
||
logger.info(f'Metadata db: {metadata_db.path} ({len(metadata_db)} entries)')
|
||
|
||
post_process = make_lumbar_post_process()
|
||
|
||
for key, value in label_map.items():
|
||
data_dir = os.path.join(data_root, key)
|
||
label_dir = os.path.join(label_root, value)
|
||
process_dataset(data_dir, label_dir, output_dir, max_images=args.max_images,
|
||
post_process=post_process, max_z_spacing=MAX_Z_SPACING_MM,
|
||
allowed_levels=LUMBAR_LEVELS, min_levels=MIN_LUMBAR_LEVELS,
|
||
metadata_cache=metadata_db)
|
||
|
||
if __name__ == '__main__':
|
||
main() |