CBT_project/imaging/preprocessing.py

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.")