CBT_project/xfr_check_spinous.py

#!/home/xfr/.conda/envs/cbt/bin/python
"""
檢查各 (volume, level) 的棘突（midline 後側構造）是否缺如
（先前手術如 laminectomy / 棘突切除，例：0005 L5、0770 L5）。

判定原理（見 imaging.orientation.diagnose_spinous_process）：
  棘突完整時，尖端是全椎體最後側的骨且位於中線 →
    deficit（全骨最後側 - 中線帶最後側）≈ 0，後側中線窄帶有骨。
  切除後，最後側骨偏到側方殘餘、後側中線空洞 → deficit 大 + rear3 ≈ 0。
  no_spinous = deficit >= 4 vox 且 rear3 <= 20 vox（正常 level deficit 恒為 0）。

兩種模式：
  quick（預設，掃全資料集用）→ 兩段式：
    stage 1：axial 投影快速篩檢（~0.5s/level）。用「寬行（椎體行）」擬合
          中線軌跡（對傾斜/旋轉穩健，不受後側殘片污染），找出「缺棘突候選」。
          高召回（幾乎不漏）、但因純投影在傾斜脊椎上仍會有少數假候補。
    stage 2：僅對候選跑 full（鏡稱面）驗證（~3s/候選），以 full 為準給出
          最終判定（confirmed / false alarm）。全資料集 ~25min + 候數×3s。
  full（指定單一 volume+level 時預設）：跑鏡稱面 + 棘突/椎體完整分割，
          另回報 VBODY 分割 mode（nosp_gap / nosp_post_min / quantile...），
          確認該 level 的椎體切分落在「椎體後側末端（體/弓最薄處）」。

識別出的 no_spinous level 在 xfr_debug.py 執行時會打 [NO-SP] 標記，
並自動改用放寬後側谷底的椎體切分（不再把殘留後側要素當棘突移除）。

Usage:
  python xfr_check_spinous.py                    # quick，全部 volume L1~L5
  python xfr_check_spinous.py 0005               # quick，該 volume L1~L5
  python xfr_check_spinous.py 0005 L5            # full（精確）單一 level
  python xfr_check_spinous.py --full 0005 L5     # 強制 full
  python xfr_check_spinous.py --quick            # 強制 quick（全資料集掃）
"""

import csv
import os
import sys
import time

import numpy as np
import SimpleITK as sitk

from imaging.orientation import (best_symmetry_plane, segment_spinous_process,
                                 best_upper_endplate_plane, segment_vertebral_body,
                                 diagnose_spinous_process)

standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/'

LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')

LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')

# 與 diagnose_spinous_process 相同的判定閾值
DEFICIT_MIN = 4.0
REAR3_MAX = 20
# quick：中線帶內一行的最小 bone voxel 數（< 此值視為殘片游離 voxel，不算「尖端在」）
MIN_MID_VOX = 3


# =====================================================================
# quick screen：不依賴鏡稱面（~0.3s/level）
# =====================================================================
def quick_screen(m):
    """axial 投影 (y, x) 快速篩檢（不跑鏡稱面，對傾斜/旋轉穩健）：
      - 「寬行」= 行內骨 voxel 數 >= 最大行的 50%（即椎體行；椎體左右
        大致對稱 → 寬行 x 重心 ≈ 中線上的點）。
      - 中線軌跡 x(y) = p0 + p1·y 只用寬行（椎體行）擬合：
        * 用全部行擬合 → laminectomy 後側殘片行污染中線，deficit 被吃掉
          （0005 L5 假正常）；
        * 用固定垂直中線 → 傾斜棘突尖端被誤判為側方偏移
          （0555/0609/0637 假切除）。椎體行兩者都不沾，故最穩。
      - deficit = 最後側骨行 (y_min) 到「中線軌跡 ± 25% 椎體寬度 band
        內有骨的最後側行」的行距（尖端在中線上 → 0；切除後 → ≈切除量）。
      - rear3/rear6 = 最後側 3/6 行內、中線軌跡 ± 10% 寬度（窄中線帶）的
        骨 voxel 數。
      - top_off = 最後側一行 x 重心到中線軌跡的距離。
    回傳 dict 或 None（資料不足）。"""
    proj = m.max(axis=0)
    ys_p, xs_p = np.where(proj > 0)
    if ys_p.size < 50:
        return None
    ny = proj.shape[0]
    cnt = np.bincount(ys_p, minlength=ny)
    max_cnt = int(cnt.max())
    if max_cnt < 4:
        return None
    xsum = np.bincount(ys_p, weights=xs_p.astype(np.float64), minlength=ny)
    centroid = xsum / cnt.astype(float)
    wide = cnt >= (0.5 * max_cnt)          # 椎體行
    if int(wide.sum()) < 10:
        return None
    yv = np.where(wide)[0]
    p1, p0 = np.polyfit(yv, centroid[wide], 1)   # 中線軌跡（椎體行擬合）
    xmax = np.full(ny, -1.0)
    xmin = np.full(ny, np.inf)
    np.maximum.at(xmax, ys_p, xs_p)
    np.minimum.at(xmin, ys_p, xs_p)
    span_c = float(np.median(xmax[yv] - xmin[yv]))
    if span_c < 10:
        return None
    line = p0 + p1 * ys_p
    band = 0.25 * span_c                   # central band（椎體中央區）
    narrow = 0.10 * span_c                  # 窄中線帶
    y_min = int(ys_p.min())
    # 「中線有骨」需 >= MIN_MID_VOX 個 voxel 才算數：棘突尖端在帶內有
    # 多個 voxel，切除後的側方殘片在帶內通常只有 1-2 個游離 voxel
    # （否則 0005 L5 的單 voxel 殘端會被當成尖端 → 假正常）。
    bcnt = np.bincount(ys_p[np.abs(xs_p - line) <= band], minlength=ny)
    idx = np.where(bcnt >= MIN_MID_VOX)[0]
    y_mid = int(idx.min()) if idx.size else ny - 1
    deficit = float(y_mid - y_min)
    nn = np.abs(xs_p - line) <= narrow
    rear3 = int(((ys_p <= y_min + 3.0) & nn).sum())
    rear6 = int(((ys_p <= y_min + 6.0) & nn).sum())
    top_off = float(abs(xs_p[ys_p == y_min].mean() - (p0 + p1 * y_min)))
    return {'deficit': deficit, 'rear3': rear3, 'rear6': rear6,
            'top_off': top_off, 'no_spinous': bool(deficit >= DEFICIT_MIN
                                                   and rear3 <= REAR3_MAX)}


# =====================================================================
# full：鏡稱面 + 完整棘突/椎體分割（精確，較慢）
# =====================================================================
def full_check(m):
    sym = best_symmetry_plane(m)
    diag = diagnose_spinous_process(m, sym)
    sp_mask, sp_th, sp_info = segment_spinous_process(m, sym)
    ep = best_upper_endplate_plane(m)
    vb, vb_th, vb_info = segment_vertebral_body(m, sym, ep, sp_th, sp_info['mode'])
    row = {'deficit': float(diag['deficit']), 'rear3': diag['rear3'],
           'rear6': diag['rear6'], 'top_off': diag['top_off'],
           'no_spinous': bool(diag['no_spinous']),
           'sp_mode': sp_info['mode'],
           'sp_pct': 100.0 * sp_info['n_sp'] / max(int(m.sum()), 1),
           'vb_mode': vb_info['mode'],
           'vb_pct': (100.0 * vb_info['n_vb'] / max(int(m.sum()), 1))
                     if vb is not None else None,
           'vb_ap_th': vb_th}
    return row


# =====================================================================
def load_binary(volume_dir, level):
    # _binary.nii.gz 現為原解析度（未插值）；0.5mm 分析優 _binary_sdf（SDF 平滑遮罩）
    p = os.path.join(volume_dir, f'{level}_binary_sdf.nii.gz')
    if not os.path.exists(p):
        p = os.path.join(volume_dir, f'{level}_binary.nii.gz')
    if not os.path.exists(p):
        return None
    m = sitk.GetArrayFromImage(sitk.ReadImage(p, sitk.sitkUInt8)) > 0
    if int(m.sum()) < 1000 or min(m.shape) < 8:
        return None
    return m


def check_level(vid, level, use_full):
    m = load_binary(os.path.join(standardized_dir, vid), level)
    if m is None:
        return {'vid': vid, 'level': level, 'note': 'missing/degenerate mask'}
    t0 = time.time()
    row = quick_screen(m) if not use_full else full_check(m)
    row.update(vid=vid.rsplit('.', 1)[-1], level=level,
               n_bone=int(m.sum()), sec=round(time.time() - t0, 2))
    return row


def parse_args(argv):
    use_full = None
    pos = []
    for a in argv:
        if a == '--full':
            use_full = True
        elif a == '--quick':
            use_full = False
        else:
            pos.append(a)
    if len(pos) > 2:
        sys.exit('Usage: python xfr_check_spinous.py [--full|--quick] [volume_id] [level]')
    vid_arg = pos[0] if len(pos) >= 1 else None
    level_arg = pos[1] if len(pos) >= 2 else None
    if level_arg and level_arg.upper() not in LEVELS:
        sys.exit(f'Invalid level: {level_arg} (choose from {"/".join(LEVELS)})')
    if level_arg and not vid_arg:
        sys.exit('level requires volume_id')
    if use_full is None:
        use_full = bool(vid_arg and level_arg)     # 指定單 level → 精確模式
    return vid_arg, (level_arg.upper() if level_arg else None), use_full


def main():
    os.makedirs(LOG_DIR, exist_ok=True)
    vid_arg, level_arg, use_full = parse_args(sys.argv[1:])

    volumes = [d for d in sorted(os.listdir(standardized_dir))
               if os.path.isdir(os.path.join(standardized_dir, d))]
    if vid_arg is not None:
        key = vid_arg.lower()
        vols = [v for v in volumes
                if v.lower() == key or v.rsplit('.', 1)[-1] == key]
        if not vols:
            sys.exit(f'Volume not found: {vid_arg}')
        volumes = vols
    levels = (level_arg,) if level_arg else LEVELS
    tasks = [(v, lvl) for v in volumes for lvl in levels]
    mode = ('FULL (mirror plane)' if use_full
        else 'QUICK screen -> FULL-verify candidates (2-stage)')
    print(f'Mode: {mode} | {len(volumes)} volumes x {len(levels)} levels = {len(tasks)} tasks',
          flush=True)

    t0 = time.time()
    rows = []
    candidates = []      # (row 索引, volume 目錄名)
    for i, (vid, lvl) in enumerate(tasks, 1):
        try:
            row = check_level(vid, lvl, use_full)
        except Exception as e:
            row = {'vid': vid.rsplit('.', 1)[-1], 'level': lvl, 'note': f'error: {e}'}
        rows.append(row)
        if 'note' in row:
            print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: {row['note']}", flush=True)
        elif use_full:
            top_str = 'n/a' if row['top_off'] is None else f"{row['top_off']:.1f}"
            print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: "
                  f"deficit={row['deficit']:6.1f}  rear3={row['rear3']}  "
                  f"top_off={top_str}  {'<<< NO-SPINUS' if row['no_spinous'] else 'ok'}  "
                  f"sp={row['sp_mode']} {row['sp_pct']:.1f}%  "
                  f"vb={row['vb_mode']} {row['vb_pct']:.1f}%", flush=True)
        else:
            mark = 'candidate' if row['no_spinous'] else 'ok'
            print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: "
                  f"deficit={row['deficit']:6.1f}  rear3={row['rear3']:4d}  "
                  f"rear6={row['rear6']:5d}  top_off={row['top_off']:5.1f}  {mark}",
                  flush=True)
        if not use_full and row.get('no_spinous'):
            candidates.append((len(rows) - 1, vid))

    # ---- Stage 2（quick 模式）：候選用 full（鏡稱面）驗證，以 full 為準 ----
    flagged = []
    cleared = []
    if use_full:
        for r in rows:
            if r.get('no_spinous'):
                flagged.append(f"{r['vid']} {r['level']}")
    elif candidates:
        print(f'\n=== stage 2: {len(candidates)} quick candidate(s) -> '
              f'full (mirror plane) verification ===', flush=True)
        for j, (k, vid_dir) in enumerate(candidates, 1):
            row = rows[k]
            try:
                m = load_binary(os.path.join(standardized_dir, vid_dir), row['level'])
                frow = full_check(m) if m is not None else {}
            except Exception as e:
                frow = {'note': f'verify error: {e}'}
            for kk in ('deficit', 'rear3', 'rear6', 'top_off', 'no_spinous',
                       'sp_mode', 'sp_pct', 'vb_mode', 'vb_pct', 'vb_ap_th'):
                if kk in frow:
                    row[kk] = frow[kk]
            if not frow:
                flag_str = '-> UNVERIFIED (mask missing/full failed)'
                cleared.append(f"{row['vid']} {row['level']} (UNVERIFIED)")
            elif row.get('no_spinous'):
                flag_str = '<<< NO-SPINUS (confirmed)'
                flagged.append(f"{row['vid']} {row['level']}")
            else:
                flag_str = '-> normal (quick false alarm)'
                cleared.append(f"{row['vid']} {row['level']} "
                               f"(full deficit={row.get('deficit', float('nan')):.1f})")
            to = row.get('top_off')
            print(f"[verify {j}/{len(candidates)}] {row['vid']} {row['level']}: "
                  f"deficit={row.get('deficit', float('nan')):6.1f}  "
                  f"rear3={row.get('rear3', '?')}  "
                  f"top_off={'n/a' if to is None else format(to, '5.1f')}  "
                  f"sp={row.get('sp_mode', '?')}  "
                  f"vb={row.get('vb_mode', '?')} {row.get('vb_pct') or 0:.1f}%  {flag_str}",
                  flush=True)

    # ---- CSV ----
    ts = time.strftime('%Y%m%d_%H%M%S')
    fields = ['vid', 'level', 'n_bone', 'deficit', 'rear3', 'rear6', 'top_off',
              'no_spinous', 'sp_mode', 'sp_pct', 'vb_mode', 'vb_pct', 'vb_ap_th',
              'sec', 'note']
    csv_path = os.path.join(LOG_DIR, f'spinous_check_{ts}.csv')
    with open(csv_path, 'w', newline='') as f:
        wtr = csv.DictWriter(f, fieldnames=fields, extrasaction='ignore')
        wtr.writeheader()
        for r in rows:
            r['vid'] = r.get('vid', '')
            wtr.writerow(r)

    dt = (time.time() - t0) / 60.0
    print('\n' + '=' * 64)
    print(f'Done in {dt:.1f} min | {len(rows)} levels checked (mode: {mode})')
    if flagged:
        print(f'\n>>> {len(flagged)} level(s) with NO SPINOUS PROCESS (prior laminectomy/resection):')
        for f_ in flagged:
            print(f'    {f_}')
        print('    (xfr_debug.py 執行這些 level 時會打 [NO-SP] 並自動改用放寬後側谷底的椎體切分)')
    else:
        print('\nNo missing spinous process detected.')
    if cleared:
        print(f'    (quick 候選、經 full 驗證為正常: {", ".join(cleared)})')
    if use_full:
        quant = [f"{r['vid']} {r['level']}" for r in rows
                 if r.get('vb_mode') == 'quantile']
        if quant:
            print(f'\nWARNING: VBODY 退回 55 百分位切分（可能切進椎體）: {", ".join(quant)}')
    print(f'CSV: {csv_path}')


if __name__ == '__main__':
    main()