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.
279 lines
No EOL
12 KiB
Python
279 lines
No EOL
12 KiB
Python
import os
|
||
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
|
||
|
||
|
||
|
||
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)
|
||
|
||
# 建立每個檔案的輸出資料夾
|
||
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.") |