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.
424 lines
No EOL
20 KiB
Python
424 lines
No EOL
20 KiB
Python
#!/home/xfr/.conda/envs/cbt/bin/python
|
||
"""把 PSO 找到的 CBT 螺絲位置（旋轉 0.5mm 標準系）映回原始
|
||
（未旋轉、未重取樣）native CT 空間，整支 volume 各 level 的螺絲
|
||
（L1-L5 x L/R，最多 10 支）存成單一 label 體積：
|
||
|
||
Output_dir/<run_date>/<volume_id>/cbt.nii.gz
|
||
Output_dir/<run_date>/<volume_id>/x-ap.jpg 合成前後(AP) X 光投影视圖
|
||
Output_dir/<run_date>/<volume_id>/x-lat.jpg 合成側位 X 光投影视圖
|
||
（只投影脊椎骨 + 螺絲、不含軟組織；
|
||
見 render_xray_projections）
|
||
|
||
label 值：L1L=1 L1R=2 L2L=3 L2R=4 L3L=5 L3R=6 L4L=7 L4R=8 L5L=9 L5R=10
|
||
（0 = 背景）。
|
||
|
||
座標鏈首選：直接用預處理寫檔時記錄在 <vol>/transform.json 的正向鏈
|
||
（imaging.transforms.original_to_source：boxes / std_flip_axes / ap_flip /
|
||
rotated R、center、start / 原 CT 幾何；純 index 空間、不依賴 standardized
|
||
輸出的物理 header）。該鏈為仿射（original index <-> rotated disk index），
|
||
螺絲點經仿射逆向 o = M^-1 (s - t) 精確映射——不重新估算任何平面 / 幾何、
|
||
不重取樣 native label。
|
||
|
||
fallback（僅舊世代 volume 沒有 transform.json、或該 level 沒有 rotated
|
||
記錄時）：純 index 空間重算 R/center（template 骨頭 mask 的
|
||
best_symmetry_plane / best_upper_endplate_plane）、fstart（rotated 檔
|
||
origin 反映到 template index 的整數）、bbox2（native label 0.5mm 線性
|
||
重取樣 >0.5 的最大 26-連通區域 bbox）：
|
||
rotated disk r -> t = R^T (r + fstart - c) + c
|
||
-> g0 = bbox2s + (wx-1-tx, wy-1-ty, tz)；ap_flip 時 g0y = N05y-1-g0y
|
||
-> native: rint(g0 * 0.5 / sn)
|
||
重算結果可能與生成時參數不完全一致：2026-09-09 驗證 liver_100 螺絲柱
|
||
in-label 由 93-99%（transform.json）掉到 30-94%（重算）；新世代資料
|
||
（transform.json 在）一律走主路徑。
|
||
|
||
螺絲參數化（rotated 系、0.5mm index）：pos = [z, y, x, az°, alt°, d mm, L mm]；
|
||
方向 d_v = (cos az sin alt, sin az sin alt, cos alt)；末端 = p0 + L/0.5 * d_v；
|
||
柱半徑 = d/2 mm。
|
||
|
||
驗證（2026-09-09，transform.json 路徑）：20260909_141448 run —
|
||
0001（10/10）、liver_100（10/10）、covid19（L1/L2 六側）螺絲柱體
|
||
in-label >= 92%、2-voxel 膨脹 >= 98%。covid19 L3 兩側除外：該 volume 的
|
||
L3 旋轉 bone mask 是 4.5k voxel / 5 slice 的退化區域（L4/L5 預處理缺檔、
|
||
先前棘突切除），屬資料問題非映射問題（L3 bone 碎片本身逆向 in-label 97.8%）。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from tinydb import TinyDB, Query
|
||
|
||
from config.constant import LABEL_MAP
|
||
from imaging.resample import resample_img
|
||
from imaging.segmentation import _largest_cc_bbox
|
||
from imaging.orientation import best_symmetry_plane, best_upper_endplate_plane
|
||
from imaging.transforms import (level_file_path, load_transform,
|
||
original_to_source, transform_path)
|
||
from visualization.res_bone_figure import compute_normalizing_rotation
|
||
|
||
standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3'
|
||
data_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
|
||
label_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/'
|
||
SUBDIRS = {
|
||
'colon': 'conlon',
|
||
'COVID-19': 'COVID-19',
|
||
'HNSCC-3DCT-RT_neck': 'HNSCC-3DCT-RT_neck',
|
||
'liver': 'Liver',
|
||
}
|
||
Output_dir = '/mnt/1248/open/cyrou/Output'
|
||
|
||
_PROJ_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
LOG_DIR = os.path.join(_PROJ_DIR, 'logs')
|
||
SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results')
|
||
META_DB = os.path.join(_PROJ_DIR, 'xfr_image_metadata.json')
|
||
|
||
LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')
|
||
LEVEL_LABEL_VAL = {v: int(k) for k, v in LABEL_MAP.items() if v in LEVELS} # {'L1': 20, ...}
|
||
PROJ_FILENAME = {'ap': 'x-ap.jpg', 'lateral': 'x-lat.jpg'}
|
||
|
||
logger = logging.getLogger('xfr_cbt_native')
|
||
|
||
|
||
def find_native_paths(volume_id):
|
||
"""回傳 (ct_path, label_path, ap_flip)；找不到 native 檔對時 raise。"""
|
||
ap_flip = False
|
||
try:
|
||
meta = TinyDB(META_DB, access_mode='r').table('images') \
|
||
.get(Query().name == volume_id)
|
||
if meta is None:
|
||
logger.warning(f'{volume_id}: not in metadata db; ap_flip defaults to False')
|
||
else:
|
||
ap_flip = bool(meta.get('ap_flip', False))
|
||
except Exception as e:
|
||
logger.warning(f'{volume_id}: metadata db unavailable ({e}); ap_flip defaults to False')
|
||
for sub, lab_sub in SUBDIRS.items():
|
||
ct = f'{data_root}{sub}/{volume_id}.nii.gz'
|
||
lb = f'{label_root}{lab_sub}/{volume_id}_seg.nii.gz'
|
||
if os.path.isfile(ct) and os.path.isfile(lb):
|
||
return ct, lb, ap_flip
|
||
raise FileNotFoundError(f'{volume_id}: no native CT/label pair under {data_root}')
|
||
|
||
|
||
def level_geometry(volume_id, level, lb_img, lb_arr):
|
||
"""每 level 的 (R, c, fstart, bbox2s, wx, wy)；與預處理/驗證腳本同定義。"""
|
||
vol_dir = os.path.join(standardized_dir, volume_id)
|
||
# 未旋轉檔：新世代在 crop/ 子資料夾、舊世代在頂層
|
||
smd_img = sitk.ReadImage(level_file_path(vol_dir, level, 'smd_resampled'))
|
||
smd = sitk.GetArrayFromImage(smd_img).astype(np.float32)
|
||
if not (np.isfinite(smd).all() and smd.min() < 0 and smd.max() > 0):
|
||
b = (sitk.GetArrayFromImage(
|
||
sitk.ReadImage(level_file_path(vol_dir, level, 'binary_nn'))) > 0).astype(np.uint8)
|
||
else:
|
||
b = (smd < 0.5).astype(np.uint8)
|
||
sym = best_symmetry_plane(b)
|
||
symp = best_upper_endplate_plane(b)
|
||
R, c = compute_normalizing_rotation(b, sym, symp)
|
||
rot_img = sitk.ReadImage(f'{vol_dir}/rotated/{level}_binary_sdf.nii.gz')
|
||
fstart = np.round(np.array(
|
||
smd_img.TransformPhysicalPointToIndex(rot_img.GetOrigin()))).astype(float)
|
||
|
||
lv = LEVEL_LABEL_VAL[level]
|
||
bin_img = sitk.GetImageFromArray((lb_arr == lv).astype(np.uint8))
|
||
bin_img.CopyInformation(lb_img)
|
||
bin_lin = resample_img(sitk.Cast(bin_img, sitk.sitkFloat32))
|
||
m_full = sitk.GetImageFromArray((sitk.GetArrayFromImage(bin_lin) > 0.5).astype(np.uint8))
|
||
m_full.CopyInformation(bin_lin)
|
||
cc = _largest_cc_bbox(m_full)
|
||
if cc is None:
|
||
raise ValueError(f'{volume_id} {level}: empty native level mask (label {lv})')
|
||
_, bbox2 = cc
|
||
wx, wy = smd_img.GetSize()[0], smd_img.GetSize()[1]
|
||
return R, c, fstart, np.array(bbox2[:3], float), wx, wy
|
||
|
||
|
||
def rotated_to_native(r_xyz, geom, N05, sn, ap_flip):
|
||
"""r_xyz: (N,3) rotated disk (x,y,z) index -> (N,3) native 連續 index。"""
|
||
R, c, fstart, bbox2s, wx, wy = geom
|
||
t = (np.asarray(r_xyz, float) + fstart - c) @ R + c
|
||
g0 = np.empty_like(t)
|
||
g0[:, 0] = bbox2s[0] + (wx - 1 - t[:, 0])
|
||
g0[:, 1] = bbox2s[1] + (wy - 1 - t[:, 1])
|
||
g0[:, 2] = bbox2s[2] + t[:, 2]
|
||
if ap_flip:
|
||
g0[:, 1] = N05[1] - 1 - g0[:, 1]
|
||
return g0 * (0.5 / sn)
|
||
|
||
|
||
def screw_voxel_xyz(pos):
|
||
"""pos = [z, y, x, az°, alt°, d, L]（rotated 系 0.5mm index）。
|
||
回傳柱體內 (x,y,z) float voxel 座標（N,3）。"""
|
||
z, y, x, az, alt, d, L = (float(v) for v in pos[:7])
|
||
azr, altr = np.radians(az), np.radians(alt)
|
||
ca, sa = np.cos(azr), np.sin(azr)
|
||
ct_, st_ = np.cos(altr), np.sin(altr)
|
||
dv = np.array([ca * st_, sa * st_, ct_])
|
||
e1 = np.array([ca * ct_, sa * ct_, -st_])
|
||
e2 = np.array([-sa, ca, 0.0])
|
||
p0 = np.array([x, y, z])
|
||
p1 = p0 + (L / 0.5) * dv
|
||
rad = d / 0.5 + 3.0
|
||
lo = np.floor(np.minimum(p0, p1)) - rad
|
||
hi = np.ceil(np.maximum(p0, p1)) + rad
|
||
xs, ys, zs = np.meshgrid(np.arange(lo[0], hi[0] + 1),
|
||
np.arange(lo[1], hi[1] + 1),
|
||
np.arange(lo[2], hi[2] + 1), indexing='xy')
|
||
P = np.stack([xs.ravel(), ys.ravel(), zs.ravel()], 1)
|
||
dP = P - p0
|
||
xr, yr, zr = dP @ e1, dP @ e2, dP @ dv
|
||
rrad = d # (d/2 mm 半徑) / 0.5mm 每 voxel = d 個 voxel
|
||
m = (xr ** 2 + yr ** 2 <= rrad ** 2) & (zr >= 0) & (zr <= L / 0.5)
|
||
return P[m]
|
||
|
||
|
||
def write_volume_cbt(volume_id, run_id, output_root=Output_dir, date=None):
|
||
"""把一個 volume 的 side_results 螺絲全部映回 native 空間寫 cbt.nii.gz。
|
||
回傳 (path, n_screws)；該 run 無此 volume 時回傳 (None, 0)。"""
|
||
date = date or run_id[:8]
|
||
side_vol_dir = os.path.join(SIDE_RESULT_DIR, run_id, volume_id)
|
||
if not os.path.isdir(side_vol_dir):
|
||
logger.warning(f'{volume_id}: no side results under {side_vol_dir}; cbt skipped')
|
||
return None, 0
|
||
if not any(os.path.isfile(os.path.join(side_vol_dir, f'{l}_{s}.json'))
|
||
for l in LEVELS for s in ('L', 'R')):
|
||
logger.warning(f'{volume_id}: no <level>_<side>.json in {side_vol_dir}; cbt skipped')
|
||
return None, 0
|
||
|
||
ct_path, lb_path, ap_flip = find_native_paths(volume_id)
|
||
ct = sitk.ReadImage(ct_path)
|
||
n_xyz = np.array(ct.GetSize(), int) # (x,y,z)
|
||
n_zyx = n_xyz[::-1].copy() # 輸出 (z,y,x)
|
||
sn = np.array(ct.GetSpacing(), float)
|
||
N05 = np.maximum(1, np.ceil(n_xyz.astype(float) * sn / 0.5 - 1e-6).astype(int))
|
||
|
||
# 主路徑：transform.json 記錄的正向鏈（仿射），逆向 o = M^-1 (s - t)
|
||
meta = None
|
||
try:
|
||
meta = load_transform(os.path.join(standardized_dir, volume_id))
|
||
except FileNotFoundError:
|
||
logger.warning(f'{volume_id}: no transform.json '
|
||
f'({transform_path(os.path.join(standardized_dir, volume_id))}); '
|
||
f'全部 level 退回重算路徑（近似）')
|
||
if meta is not None and not np.array_equal(
|
||
np.asarray(meta['original']['size'], int), n_xyz):
|
||
logger.warning(f'{volume_id}: transform.json original size '
|
||
f'{meta["original"]["size"]} != CT {list(n_xyz)}; '
|
||
f'忽略 transform.json，退回重算路徑')
|
||
meta = None
|
||
|
||
out_arr = np.zeros(n_zyx, np.uint8)
|
||
lb_img = lb_arr = None # fallback 才需要讀 native label
|
||
geom_cache = {} # fallback：每 level 的 (R, c, fstart, bbox2s, wx, wy)
|
||
aff_cache = {} # 主路徑：每 level 的 (M^-1, t)
|
||
n_screws, skipped = 0, []
|
||
n_tf = n_fb = 0
|
||
for li, level in enumerate(LEVELS):
|
||
for side in ('L', 'R'):
|
||
jp = os.path.join(side_vol_dir, f'{level}_{side}.json')
|
||
if not os.path.isfile(jp):
|
||
continue
|
||
pos = json.load(open(jp))['position']
|
||
try:
|
||
cyl = screw_voxel_xyz(pos)
|
||
lv = meta['levels'].get(level) if meta is not None else None
|
||
if lv is not None and 'rotated' in lv:
|
||
if level not in aff_cache:
|
||
M, t = original_to_source(meta, level, 'rotated')
|
||
aff_cache[level] = (np.linalg.inv(M), t)
|
||
Mi, t = aff_cache[level]
|
||
nat = (cyl - t) @ Mi.T # (N,3) 原 CT 連續 index
|
||
n_tf += 1
|
||
else:
|
||
if lb_arr is None:
|
||
lb_img = sitk.ReadImage(lb_path)
|
||
lb_arr = sitk.GetArrayFromImage(lb_img)
|
||
if level not in geom_cache:
|
||
geom_cache[level] = level_geometry(volume_id, level, lb_img, lb_arr)
|
||
nat = rotated_to_native(cyl, geom_cache[level], N05, sn, ap_flip)
|
||
n_fb += 1
|
||
ni = np.rint(nat).astype(int)
|
||
valid = (ni >= 0).all(1) & (ni < n_xyz).all(1)
|
||
idx = ni[valid]
|
||
if idx.shape[0] == 0:
|
||
raise ValueError('all screw voxels out of original CT bounds')
|
||
val = li * 2 + 1 + (1 if side == 'R' else 0)
|
||
out_arr[idx[:, 2], idx[:, 1], idx[:, 0]] = val
|
||
n_screws += 1
|
||
except Exception as e:
|
||
logger.error(f'{volume_id} {level}_{side}: {e}')
|
||
skipped.append(f'{level}_{side} ({e})')
|
||
|
||
if n_screws == 0:
|
||
raise ValueError(f'{volume_id}: no screw could be written (skipped: {skipped or "none read"})')
|
||
out_dir = os.path.join(output_root, date, volume_id)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
out_img = sitk.GetImageFromArray(out_arr)
|
||
out_img.CopyInformation(ct)
|
||
out_path = os.path.join(out_dir, 'cbt.nii.gz')
|
||
sitk.WriteImage(out_img, out_path)
|
||
if skipped:
|
||
logger.warning(f'{volume_id}: failed sides: {", ".join(skipped)}')
|
||
logger.info(f'{volume_id}: {n_screws}/10 screws -> {out_path} '
|
||
f'(transform.json={n_tf}, re-est={n_fb}, ap_flip={ap_flip})')
|
||
|
||
# 收尾：原 CT + 螺絲 -> 合成 AP / Lateral X 光投影视圖 (x-ap.jpg / x-lat.jpg)
|
||
try:
|
||
render_xray_projections(out_path, ct_path, out_dir)
|
||
except Exception as e:
|
||
logger.warning(f'{volume_id}: X-ray projection render failed: {e}')
|
||
return out_path, n_screws
|
||
|
||
|
||
def render_xray_projections(cbt_path, ct_path, out_dir,
|
||
views=('ap', 'lateral'), margin_mm=50.0):
|
||
"""原 CT + cbt.nii.gz（native 同一 grid 的螺絲 label）-> 合成「只有骨頭」X 光投影视圖。
|
||
|
||
只投影脊椎骨 + 螺絲，不含軟組織：
|
||
- spine mask：native 分割 label（1-24 = C1..L5，見 config.constant.LABEL_MAP）；
|
||
label 缺 / grid 不符時退回 HU 300-3000 閾值。
|
||
- 骨 μ = clip(HU, 0, 2000)/400（松質骨~0.2-1、皮質/終板~2-5）。
|
||
- 螺絲 voxel μ = 60（金屬等效，最亮白）。
|
||
投影视圖 = 沿中心射線 μ 的線積分（AP 沿 y、Lateral 沿 x）；裁到螺絲 bbox 外
|
||
margin_mm 的 spine 區域；骨用 percentile(1,99) 獨立視窗，螺絲再疊加裁白。
|
||
方向：SimpleITK LPS（x→右、y→後、z→上）；AP 頂=上、病人右在畫面左（R 標記）；
|
||
Lateral 頂=上、前位在左、後位在右（A/P 標記）。
|
||
spacing 各軸不等時先重取樣到 min(spacing) 各向同性格，維持投影 aspect ratio。
|
||
輸出 x-ap.jpg / x-lat.jpg，回傳 {view: jpg_path}。"""
|
||
import matplotlib
|
||
matplotlib.use('Agg')
|
||
import matplotlib.pyplot as plt
|
||
|
||
vid = os.path.basename(os.path.normpath(out_dir))
|
||
ct = sitk.ReadImage(ct_path)
|
||
cbt = sitk.ReadImage(cbt_path)
|
||
if ct.GetSize() != cbt.GetSize():
|
||
raise ValueError(f'cbt/CT grid 尺寸不符 {cbt.GetSize()} vs {ct.GetSize()}')
|
||
sp = np.array(ct.GetSpacing(), float) # (x,y,z)
|
||
native_size = ct.GetSize()
|
||
# spacing 各軸不等：先重取樣 CT/螺絲（與後面的分割 label）到
|
||
# min(spacing) 的各向同性格，維持投影视圖正確的 aspect ratio
|
||
min_sp = float(sp.min())
|
||
isosp = isosize = None
|
||
if not np.allclose(sp, min_sp):
|
||
isosp = [min_sp, min_sp, min_sp]
|
||
isosize = [max(1, int(round(s * v / min_sp))) for s, v in zip(sp, ct.GetSize())]
|
||
o, d = ct.GetOrigin(), ct.GetDirection()
|
||
ct = sitk.Resample(sitk.Cast(ct, sitk.sitkFloat32), isosize, sitk.Transform(),
|
||
sitk.sitkLinear, o, isosp, d, 0.0)
|
||
cbt = sitk.Resample(cbt, isosize, sitk.Transform(),
|
||
sitk.sitkNearestNeighbor, o, isosp, d, 0)
|
||
ct_arr = sitk.GetArrayFromImage(ct).astype(np.float32) # (z,y,x)
|
||
cbt_arr = sitk.GetArrayFromImage(cbt)
|
||
sp = np.array(ct.GetSpacing(), float)
|
||
|
||
metal = cbt_arr > 0
|
||
try:
|
||
_, lb_path, _ = find_native_paths(vid)
|
||
lb_img = sitk.ReadImage(lb_path)
|
||
# 比對「重取樣前」的 native grid（label 與原 CT 同格）
|
||
if lb_img.GetSize() != native_size:
|
||
raise ValueError(f'label/CT grid 尺寸不符 {lb_img.GetSize()} vs {native_size}')
|
||
if isosp is not None:
|
||
lb_img = sitk.Resample(lb_img, isosize, sitk.Transform(),
|
||
sitk.sitkNearestNeighbor,
|
||
ct.GetOrigin(), isosp, ct.GetDirection(), 0)
|
||
lb = sitk.GetArrayFromImage(lb_img)
|
||
spine = (lb >= 1) & (lb <= 24) # LABEL_MAP: C1..L5 全為脊椎
|
||
except Exception as e:
|
||
logger.warning(f'{vid}: native 分割 label 不可用（{e}）；'
|
||
f'退回 HU 300-3000 閾值當脊椎骨')
|
||
spine = (ct_arr >= 300) & (ct_arr <= 3000)
|
||
|
||
bone_mu = np.where(spine & ~metal, np.clip(ct_arr, 0, 2000) / 400.0, 0.0)
|
||
metal_mu = 60.0 * metal
|
||
|
||
z, y, x = np.where(metal)
|
||
if z.size == 0:
|
||
raise ValueError(f'{vid}: cbt 無螺絲 voxel')
|
||
mz, my, mx = (int(margin_mm / s) for s in (sp[2], sp[1], sp[0]))
|
||
zs = slice(max(0, z.min() - mz), min(ct_arr.shape[0], z.max() + mz + 1))
|
||
ys = slice(max(0, y.min() - my), min(ct_arr.shape[1], y.max() + my + 1))
|
||
xs = slice(max(0, x.min() - mx), min(ct_arr.shape[2], x.max() + mx + 1))
|
||
|
||
proj = {
|
||
'ap': (bone_mu[zs, :, xs].sum(axis=1), metal_mu[zs, :, xs].sum(axis=1)), # 沿 y 積分
|
||
'lateral': (bone_mu[zs, ys, :].sum(axis=2), metal_mu[zs, ys, :].sum(axis=2)) # 沿 x 積分
|
||
}
|
||
flip_x = {'ap': True, 'lateral': False}
|
||
markers = {'ap': ('R', 'L'), 'lateral': ('A', 'P')} # 畫面左=前位(A)、右=後位(P)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
paths = {}
|
||
for view in views:
|
||
if view not in proj:
|
||
continue
|
||
b, m = proj[view]
|
||
b = b[::-1, ::-1] if flip_x[view] else b[::-1, :] # 頂=頭端
|
||
m = m[::-1, ::-1] if flip_x[view] else m[::-1, :]
|
||
bnz = b[b > 0]
|
||
if bnz.size:
|
||
lo, hi = np.percentile(bnz, 1), np.percentile(bnz, 99)
|
||
bone = np.clip((b - lo) / (hi - lo + 1e-9), 0, 1) ** 0.7
|
||
else:
|
||
bone = np.zeros_like(b)
|
||
mmax = float(m.max())
|
||
metal_img = (m / (mmax + 1e-9)) ** 0.5 if mmax > 0 else m * 0
|
||
img = np.clip(bone + metal_img, 0, 1)
|
||
ll, rl = markers[view]
|
||
h, w = img.shape
|
||
fig, ax = plt.subplots(figsize=(8.0 * w / h, 8.0), dpi=110)
|
||
ax.imshow(img, cmap='gray', interpolation='nearest')
|
||
ax.set_title(f'{view.upper()} projection (spine + screws) - {vid}',
|
||
color='white', fontsize=13)
|
||
ax.text(0.01, 0.98, ll, transform=ax.transAxes, color='cyan', fontsize=13,
|
||
va='top', ha='left', fontweight='bold')
|
||
ax.text(0.99, 0.98, rl, transform=ax.transAxes, color='cyan', fontsize=13,
|
||
va='top', ha='right', fontweight='bold')
|
||
ax.axis('off')
|
||
fig.tight_layout(pad=0.5)
|
||
p = os.path.join(out_dir, PROJ_FILENAME[view])
|
||
fig.savefig(p, format='jpg', facecolor='black')
|
||
plt.close(fig)
|
||
paths[view] = p
|
||
logger.info(f'{vid}: {view} projection -> {p}')
|
||
return paths
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='Write <Output_dir>/<date>/<volume_id>/cbt.nii.gz (screws mapped to native space) '
|
||
'from a side_results run')
|
||
parser.add_argument('run_id', help='e.g. 20260907_230652')
|
||
parser.add_argument('volumes', nargs='*',
|
||
help='volume ids (full or trailing digits); default: all in the run')
|
||
parser.add_argument('--date', default=None, help='output date dir (default: run_id[:8])')
|
||
args = parser.parse_args()
|
||
logging.basicConfig(level=logging.INFO,
|
||
format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
|
||
datefmt='%Y-%m-%d %H:%M:%S')
|
||
run_dir = os.path.join(SIDE_RESULT_DIR, args.run_id)
|
||
if not os.path.isdir(run_dir):
|
||
sys.exit(f'side results dir not found: {run_dir}')
|
||
run_vols = sorted(d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)))
|
||
if args.volumes:
|
||
vols = [v for v in run_vols
|
||
if v in args.volumes or v.rsplit('.', 1)[-1] in args.volumes]
|
||
if not vols:
|
||
sys.exit(f'none of {args.volumes} found in run {args.run_id}')
|
||
else:
|
||
vols = run_vols
|
||
ok, fail = 0, 0
|
||
for vid in vols:
|
||
try:
|
||
write_volume_cbt(vid, args.run_id, date=args.date)
|
||
ok += 1
|
||
except Exception as e:
|
||
fail += 1
|
||
logger.error(f'{vid}: {e}')
|
||
print(f'done: {ok} written, {fail} failed / {len(vols)} volume(s)')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |