CBT_project/xfr_preprocess.py

#!/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」平面圖（bone + 平面 + VBODY [金] / 棘突 [紫] 著色），
    並在旋轉體上做 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），並畫出 VBODY / 棘突著色
    #（VBODY=金、棘突=紫，與 label map 對應）；皮質著色由未旋轉 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=False, 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 + 旋轉平面圖（含 VBODY / 棘突著色）+ 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 圖（含 VBODY / 棘突著色）+ 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).')
    parser.add_argument('--output-dir', type=str, default=None, dest='output_dir',
                        help='Override the default output dir (e.g. repair run on an '
                             'older generation). Default: the module-level output_dir.')
    args = parser.parse_args()
    # local alias（避免 rebind module-level output_dir 造成 UnboundLocalError）
    out_dir = args.output_dir if args.output_dir is not None else output_dir

    # 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()}')
    logger.info(f'Output dir: {out_dir}')

    # 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, out_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()