Xiao Furen
2ae08ac2cd
Refactor the preprocessing and segmentation pipeline to handle AP orientation variations and improve anatomical boundary detection. Key changes include: - Implement automated AP orientation detection in `process_single_image` to handle prone scans by flipping CT and labels when necessary. - Enhance `segment_spinous_process` using a gap-based approach to identify the spinal canal, providing more stable thresholds for spinous process and vertebral body segmentation. - Improve optimization search space by using the vertebral body (VBODY) projection for x/z bounding box calculation instead of the whole bone. - Refactor `render_bone_figure` to unify 2D/3D visualization and support detailed anatomical coloring (VBODY, spinous process). - Update `cl_score_torch_xfr` with more robust penalty handling for out-of-bone and null-voxel regions. - Add `retry_robust` utility to handle transient NFS file system errors. - Update `xfr_preprocess.py` to include anatomical segmentation coloring in rotated level visualizations.
354 lines
No EOL
16 KiB
Python
354 lines
No EOL
16 KiB
Python
import os
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from imaging.resample import resample_img
|
||
from imaging.affine import standardize_affine
|
||
from imaging.segmentation import seg_bone
|
||
import json
|
||
import glob
|
||
from config.constant import LABEL_MAP
|
||
from imaging.nifti_io import sitk_to_nibabel, nibabel_to_sitk
|
||
from imaging.orientation import anterior_y_side
|
||
|
||
|
||
def flip_y_sitk(img):
|
||
"""index 系 y 軸（array axis 1，前後方向）翻轉：
|
||
只翻數據、spacing/origin/direction 等幾何不變，即整顆體積的前後
|
||
朝向在 index 系翻轉（y 大側 <-> y 小側）。供 supine / prone 個案
|
||
統一前後慣例（前側 = y 大側）用。"""
|
||
arr = np.flip(sitk.GetArrayFromImage(img), axis=1)
|
||
out = sitk.GetImageFromArray(arr)
|
||
out.CopyInformation(img)
|
||
return out
|
||
|
||
|
||
|
||
def load_progress(PROGRESS_FILE):
|
||
if os.path.exists(PROGRESS_FILE):
|
||
with open(PROGRESS_FILE, "r") as f:
|
||
return json.load(f)
|
||
return {}
|
||
|
||
def save_progress(progress, PROGRESS_FILE):
|
||
with open(PROGRESS_FILE, "w") as f:
|
||
json.dump(progress, f, indent=2)
|
||
|
||
def process_single_image(image_path, label_path, output_dir_base=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None):
|
||
"""metadata_cache：可選，需提供 get(name) -> dict|None 與
|
||
put(name, dict)。dict 可含 spacing=[x,y,z]、labels=[label id]。
|
||
兩者都在 db 裡時整支跳過判定不需讀影像 / label 檔。"""
|
||
|
||
file_name = os.path.basename(image_path)
|
||
name = file_name.replace(".nii.gz", "")
|
||
|
||
# pixel spacing / labels 優先用 metadata db，避免每次 run 都讀檔
|
||
meta = metadata_cache.get(name) if metadata_cache is not None else None
|
||
image = None
|
||
label = None
|
||
spacing = meta.get("spacing") if meta is not None else None
|
||
existing_labels = meta.get("labels") if meta is not None else None
|
||
|
||
if spacing is None:
|
||
image = sitk.ReadImage(image_path)
|
||
spacing = [float(v) for v in image.GetSpacing()]
|
||
if metadata_cache is not None:
|
||
metadata_cache.put(name, {"spacing": spacing})
|
||
|
||
# z spacing 過大（低解析度掃描）跳過整支 pipeline
|
||
if max_z_spacing is not None and spacing[2] > max_z_spacing:
|
||
z_spacing = spacing[2]
|
||
print(f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm, "
|
||
f"skipping pipeline for {name}")
|
||
return {
|
||
"processed_labels": [],
|
||
"missing_labels": [],
|
||
"skipped": True,
|
||
"skip_reason": f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm"
|
||
}
|
||
|
||
if existing_labels is None:
|
||
label = sitk.ReadImage(label_path)
|
||
# 取得現有 label
|
||
lssif = sitk.LabelShapeStatisticsImageFilter()
|
||
lssif.Execute(label)
|
||
existing_labels = [int(v) for v in lssif.GetLabels()] # 例如 [1,2,3,20,21]
|
||
if metadata_cache is not None:
|
||
metadata_cache.put(name, {"labels": existing_labels})
|
||
else:
|
||
print(f"Metadata db hit for {name} (spacing={spacing}, labels={existing_labels})")
|
||
|
||
print(f"Existing labels in {os.path.basename(label_path)}: {existing_labels}")
|
||
|
||
allowed_label_list = [n for n in existing_labels
|
||
if n in LABEL_MAP and (allowed_levels is None
|
||
or LABEL_MAP[n] in allowed_levels)]
|
||
|
||
# 沒有任何符合 allowed levels（如 lumbar）的 label：不建立輸出資料夾、
|
||
# 不做重取樣，整檔跳過（不計入 max_images）
|
||
if not allowed_label_list:
|
||
print(f"No label matches allowed levels {allowed_levels} in {name}, "
|
||
f"skipping (no output folder)")
|
||
return {
|
||
"processed_labels": [],
|
||
"missing_labels": [],
|
||
"skipped": True,
|
||
"skip_reason": f"no label in allowed levels {allowed_levels}"
|
||
}
|
||
|
||
# 符合 allowed levels 的 label 少於 min_levels（如 lumbar < 2 層）：
|
||
# 整支 pipeline 跳過，不建立輸出資料夾（不計入 max_images）
|
||
if min_levels is not None and len(allowed_label_list) < min_levels:
|
||
have = ", ".join(f"{n} ({LABEL_MAP[n]})" for n in allowed_label_list)
|
||
print(f"Only {len(allowed_label_list)} allowed-level label(s) [{have}] in {name} "
|
||
f"< {min_levels}, skipping pipeline (no output folder)")
|
||
return {
|
||
"processed_labels": [],
|
||
"missing_labels": [],
|
||
"skipped": True,
|
||
"skip_reason": (f"only {len(allowed_label_list)} label(s) in allowed "
|
||
f"levels < {min_levels}")
|
||
}
|
||
|
||
# 進入 pipeline：若上面的 spacing / labels 來自 metadata db，
|
||
# 影像與 label 檔還沒讀，這裡補讀
|
||
if image is None:
|
||
image = sitk.ReadImage(image_path)
|
||
if label is None:
|
||
label = sitk.ReadImage(label_path)
|
||
|
||
# LabelStatisticsImageFilter computes statistics (e.g., mean, minimum, maximum, median) of pixel values in an image, segmented by labels in a corresponding label image.
|
||
lsif = sitk.LabelStatisticsImageFilter()
|
||
lsif.Execute(image, label)
|
||
|
||
# Assume to have some sitk image (itk_image) and label (itk_label)
|
||
resampled_sitk_img = resample_img(image, out_spacing=[0.5, 0.5, 0.5], is_label=False)
|
||
resampled_sitk_lbl = resample_img(label, out_spacing=[0.5, 0.5, 0.5], is_label=True)
|
||
|
||
# 前後（AP）方向判定：本流程最終輸出慣例是前側 = y 大側（後 = y 小側）。
|
||
# 但 prone（伏位）掃描個案經 standardize_affine 後前側落在 y 小側
|
||
# （例：CTSpine1K colon 0003、0075、0460...，全 dataset 約 1%），
|
||
# 不修正時上終板 / 棘突 / 椎體分割與 rotated/ 對齊全部反掉。
|
||
#
|
||
# 判定（對每個 allowed level 的 0.5mm label 個別做、再票決——不能
|
||
# 直接 union 多層：腰椎前凸（lordosis）下各層椎體在 AP 投影會散開，
|
||
# 椎管「空隙」被其它層的骨填掉）：
|
||
# 1) anterior_y_side 回傳工作體積（0.5mm 重取樣、standardize_affine
|
||
# 之前）grid 中椎體塊所在的 y 端（y_min / y_max / None）；
|
||
# 2) standardize_affine（nibabel 端）會在輸出 affine y 分量 < 0 時
|
||
# 再翻一次 y。注意 nibabel affine 與 SimpleITK direction 的 y 分
|
||
# 量符號相反（NIfTI RAS <-> SITK LPS），所以
|
||
# standardize_affine 會翻 y <=> direction[4] > 0；
|
||
# 最終 y 端 = pre_side（若會翻 y 則 y_min<->y_max 互換）；
|
||
# 3) 最終前側會落 y 小側時，現在先對 CT 與 label 做 y 翻轉
|
||
# （純 index 翻轉、幾何不變），與 standardize_affine 的翻轉
|
||
# 組成淨效果，使所有輸出與其它個案同方向。
|
||
# 無法判定 / 投票平手時維持原方向（寧可不翻、不誤翻）。
|
||
# 判定結果存入 metadata db（ap_flip），重跑免重算。
|
||
ap_flip = (meta or {}).get("ap_flip")
|
||
if ap_flip is None:
|
||
# standardize_affine 是否會翻轉 y（nibabel affine[1,1] < 0，
|
||
# 等价於 sitk direction[4] > 0，兩者符號相反）
|
||
std_flips_y = resampled_sitk_img.GetDirection()[4] > 0
|
||
arr_lbl = sitk.GetArrayFromImage(resampled_sitk_lbl)
|
||
votes = []
|
||
for n in allowed_label_list:
|
||
side = anterior_y_side(arr_lbl == n)
|
||
if side is not None:
|
||
votes.append(side)
|
||
n_min = votes.count("y_min")
|
||
n_max = votes.count("y_max")
|
||
if n_min > n_max:
|
||
pre_side = "y_min"
|
||
elif n_max > n_min:
|
||
pre_side = "y_max"
|
||
else:
|
||
pre_side = None
|
||
if pre_side is None:
|
||
ap_flip = False
|
||
print(f"AP orientation undetermined for {name} (votes={votes}); "
|
||
f"proceeding with default orientation (anterior = large y)")
|
||
else:
|
||
final_side = (pre_side if not std_flips_y
|
||
else ("y_min" if pre_side == "y_max" else "y_max"))
|
||
ap_flip = final_side == "y_min"
|
||
print(f"AP orientation for {name}: pre={pre_side} "
|
||
f"(std_flips_y={std_flips_y}) -> final={final_side} "
|
||
f"[votes={votes}], flip={ap_flip}")
|
||
if metadata_cache is not None:
|
||
metadata_cache.put(name, {"ap_flip": bool(ap_flip)})
|
||
if ap_flip:
|
||
# label（原解析度 raw label）也要翻：seg_bone 的主遮罩鏈
|
||
#（_binary / SMD / _binary_sdf）是用 original_label（= label）
|
||
# 算的，不是用 0.5mm resampled label；漏翻時翻轉不生效。
|
||
label = flip_y_sitk(label)
|
||
resampled_sitk_img = flip_y_sitk(resampled_sitk_img)
|
||
resampled_sitk_lbl = flip_y_sitk(resampled_sitk_lbl)
|
||
print(f"AP orientation corrected for {name}: CT and label flipped "
|
||
f"along y; outputs unified to anterior = large y")
|
||
|
||
# 建立每個檔案的輸出資料夾
|
||
file_name = os.path.basename(image_path)
|
||
name = file_name.replace(".nii.gz", "")
|
||
output_dir = os.path.join(output_dir_base, name)
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 存現有 label 到 txt
|
||
txt_path = os.path.join(output_dir, f"{name}_labels.txt")
|
||
with open(txt_path, "w") as f:
|
||
for lab in existing_labels:
|
||
f.write(f"{lab}\t{LABEL_MAP.get(lab, 'Unknown')}\n")
|
||
|
||
# 遍歷現有 label 做分割；個別 label 失敗（label_map 缺該 label、
|
||
# 或 seg_bone 報錯）只跳過該 label，不中斷整檔處理
|
||
processed = []
|
||
skipped = []
|
||
for n in existing_labels:
|
||
if n not in LABEL_MAP:
|
||
print(f"Label {n} not found in label_map, skipping this label (file continues).")
|
||
skipped.append(n)
|
||
continue
|
||
if allowed_levels is not None and LABEL_MAP[n] not in allowed_levels:
|
||
print(f"Label {n} ({LABEL_MAP[n]}) not in allowed levels, "
|
||
f"skipping this label (file continues).")
|
||
skipped.append(n)
|
||
continue
|
||
try:
|
||
res = seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_dir,
|
||
label_map=LABEL_MAP, original_label=label)
|
||
if res is None:
|
||
print(f"Label {n}: empty after largest-CC extraction, skipping this label.")
|
||
skipped.append(n)
|
||
continue
|
||
(roi_path, binary_path, roi2_path, cortical_path, binary_nn_path,
|
||
binary_linear_path, smd_path, resampled_path, binary_sdf_path,
|
||
binary_erode_path) = res
|
||
for path in [roi_path, binary_path, roi2_path, cortical_path,
|
||
binary_nn_path, binary_linear_path,
|
||
smd_path, resampled_path, binary_sdf_path,
|
||
binary_erode_path]:
|
||
if path is not None:
|
||
standardize_affine(path, output_dir)
|
||
processed.append(n)
|
||
except RuntimeError as e:
|
||
print(f"Label {n} could not be processed, skipping. Error: {e}")
|
||
skipped.append(n)
|
||
|
||
return {
|
||
"processed_labels": processed,
|
||
"missing_labels": skipped
|
||
}
|
||
|
||
def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None, max_images=None, post_process=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None):
|
||
image_files = sorted(glob.glob(os.path.join(image_dir, "*.nii.gz")))
|
||
total_files = len(image_files)
|
||
print(f"Total files: {total_files}")
|
||
|
||
PROGRESS_FILE = os.path.join(output_dir, "progress.json")
|
||
|
||
progress = load_progress(PROGRESS_FILE)
|
||
all_file_summary = []
|
||
# max_images 統計「實際進入 pipeline 的檔數」：
|
||
# z spacing 超標被跳過的檔不計入配額，繼續掃描後續檔案
|
||
processed_count = 0
|
||
|
||
for idx, image_path in enumerate(image_files, 1):
|
||
if max_images is not None and processed_count >= max_images:
|
||
print(f"Reached max_images={max_images}, stopping early.")
|
||
break
|
||
|
||
file_name = os.path.basename(image_path)
|
||
name = file_name.replace(".nii.gz", "")
|
||
label_path = os.path.join(label_dir, file_name.replace(".nii.gz", "_seg.nii.gz"))
|
||
|
||
file_summary = {
|
||
"file_name": file_name,
|
||
"current_labels": [],
|
||
"missing_labels": []
|
||
}
|
||
|
||
if progress.get(name, {}).get("finished", False):
|
||
print(f"[{idx}/{total_files}] Already finished: {file_name}")
|
||
file_summary["current_labels"] = progress[name].get("processed_labels", [])
|
||
file_summary["missing_labels"] = progress[name].get("missing_labels", [])
|
||
processed_count += 1
|
||
all_file_summary.append(file_summary)
|
||
continue
|
||
|
||
if not os.path.exists(label_path):
|
||
print(f"[{idx}/{total_files}] Warning: label not found for {file_name}")
|
||
file_summary["note"] = "Label file not found"
|
||
processed_count += 1
|
||
all_file_summary.append(file_summary)
|
||
continue
|
||
|
||
try:
|
||
result = process_single_image(image_path, label_path, output_dir_base=output_dir,
|
||
max_z_spacing=max_z_spacing,
|
||
allowed_levels=allowed_levels,
|
||
min_levels=min_levels,
|
||
metadata_cache=metadata_cache)
|
||
# print(result)
|
||
# exit()
|
||
except Exception as e:
|
||
print(f"[{idx}/{total_files}] Error processing {file_name}: {e}")
|
||
file_summary["note"] = f"Error: {e}"
|
||
processed_count += 1
|
||
all_file_summary.append(file_summary)
|
||
continue
|
||
|
||
if result.get("skipped"):
|
||
# z spacing 超標：不計入 max_images 配額，繼續掃描後續檔案
|
||
print(f"[{idx}/{total_files}] Skipped ({result['skip_reason']}): {file_name}")
|
||
file_summary["note"] = result["skip_reason"]
|
||
all_file_summary.append(file_summary)
|
||
continue
|
||
|
||
processed_count += 1
|
||
file_summary["current_labels"] = result["processed_labels"]
|
||
file_summary["missing_labels"] = result["missing_labels"]
|
||
all_file_summary.append(file_summary)
|
||
|
||
progress[name] = {
|
||
"finished": True,
|
||
"processed_labels": result["processed_labels"],
|
||
"missing_labels": result["missing_labels"]
|
||
}
|
||
save_progress(progress, PROGRESS_FILE)
|
||
|
||
count_msg = f" | processed {processed_count}/{max_images}" if max_images is not None else ""
|
||
print(f"[{idx}/{total_files}] Finished: {file_name} | "
|
||
f"Missing labels: {result['missing_labels'] or 'None'}{count_msg}")
|
||
|
||
if post_process is not None:
|
||
try:
|
||
post_process(os.path.join(output_dir, name), result["processed_labels"])
|
||
except Exception as e:
|
||
print(f"[{idx}/{total_files}] post_process error for {name}: {e}")
|
||
|
||
# --- Summary ---
|
||
summary_path = os.path.join(output_dir, "all_files_label_summary.txt")
|
||
os.makedirs(os.path.dirname(summary_path), exist_ok=True)
|
||
print(f"\nWriting summary to {summary_path}...")
|
||
|
||
with open(summary_path, "w") as f:
|
||
f.write("--- CTSpine1K Dataset Label Summary ---\n")
|
||
f.write(f"Total files processed: {total_files}\n\n")
|
||
|
||
for summary in all_file_summary:
|
||
f.write("================================================\n")
|
||
f.write(f"File: {summary['file_name']}\n")
|
||
|
||
processed_labels_str = ", ".join([str(l) for l in summary['current_labels']])
|
||
f.write(f"Labels processed: {processed_labels_str}\n")
|
||
|
||
if summary['missing_labels']:
|
||
missing_str = ", ".join([f"{l} ({LABEL_MAP.get(l, 'Unknown')})" for l in summary['missing_labels']])
|
||
f.write(f"🚨 Missing Labels: {missing_str}\n")
|
||
else:
|
||
f.write("✅ Missing Labels: None\n")
|
||
|
||
if "note" in summary:
|
||
f.write(f"Note: {summary['note']}\n")
|
||
f.write("================================================\n\n")
|
||
|
||
print("All done! Summary file created.") |