scoliosis-inference/scoliosis_arcs.py

import numpy as np

from scoliosis_geometry import angle_diff_deg

def fill_zero_gaps(sign: np.ndarray, max_gap: int) -> np.ndarray:
    if max_gap <= 0:
        return sign
    sign = sign.copy()
    n = len(sign)
    i = 0
    while i < n:
        if sign[i] != 0:
            i += 1
            continue
        j = i
        while j < n and sign[j] == 0:
            j += 1
        left = sign[i - 1] if i > 0 else 0
        right = sign[j] if j < n else 0
        if (j - i) <= max_gap and left != 0 and left == right:
            sign[i:j] = left
        i = j
    return sign

def runs_of_same_sign(sign: np.ndarray) -> list[tuple[int, int, int]]:
    "Находит непрерывные куски одинакового знака (+1 подряд, -1 подряд)"
    n = len(sign)
    out = []
    i = 0
    while i < n:
        if sign[i] == 0:
            i += 1
            continue
        sgn = int(sign[i])
        s = i
        while i < n and sign[i] == sgn:
            i += 1
        e = i - 1
        out.append((sgn, s, e))
    return out

def extend_through_zeros(sign: np.ndarray, s: int, e: int, max_extend: int) -> tuple[int, int]:
    """Если дуга заканчивается рядом с нулями,
    функция слегка "расширяет" дугу в нули,
    чтобы включить край (где dx почти 0)"""
    n = len(sign)
    k = 0
    while k < max_extend and s > 0 and sign[s - 1] == 0:
        s -= 1
        k += 1
    k = 0
    while k < max_extend and e < n - 1 and sign[e + 1] == 0:
        e += 1
        k += 1
    return s, e

def dx_to_sign(dx: np.ndarray, eps: float) -> np.ndarray:
    return np.where(dx > eps, 1, np.where(dx < -eps, -1, 0)).astype(np.int8)

def compute_eps(dx: np.ndarray, eps_px: float, eps_rel: float = 0.01) -> float:
    maxabs = float(np.max(np.abs(dx)) + 1e-6)
    return max(float(eps_px), float(eps_rel) * maxabs)

def detect_arcs_from_dx(dx, eps_px=2.0, min_len=3, zero_gap=1, edge_zero_extend=2):
    "Основная детекция дуг"
    dx = np.asarray(dx, dtype=np.float32)
    if len(dx) < 2:
        return []

    eps = compute_eps(dx, eps_px, eps_rel=0.01)
    sign = dx_to_sign(dx, eps)
    sign = fill_zero_gaps(sign, zero_gap)

    arcs = []
    for sgn, s, e in runs_of_same_sign(sign):
        if (e - s + 1) >= int(min_len):
            s, e = extend_through_zeros(sign, s, e, edge_zero_extend)
            arcs.append((s, e))
    return arcs

def arc_sign(signal: np.ndarray, s: int, e: int) -> int:
    """Определяет знак дуги по сигналу на участке (берёт медиану).
    Нужно чтобы знак не ломался из-за шума"""
    return 1 if float(np.median(signal[s:e + 1])) > 0 else -1

def merge_arcs(arcs, sign_signal, max_gap=0):
    """Сливает дуги, если:
    - знак одинаковый
    - и они либо пересекаются, либо рядом (max_gap позволяет "почти рядом")
    Это финальная "склейка дуг", чтобы не было дробления"""
    if not arcs:
        return []

    arcs = sorted(arcs, key=lambda t: (t[0], t[1]))
    out = [[arcs[0][0], arcs[0][1], arc_sign(sign_signal, arcs[0][0], arcs[0][1])]]

    for s, e in arcs[1:]:
        sgn = arc_sign(sign_signal, s, e)
        ps, pe, psgn = out[-1]

        if sgn == psgn and s <= pe + max_gap:
            out[-1][1] = max(pe, e)
        else:
            out.append([s, e, sgn])

    return [(s, e) for s, e, _ in out]

def cobb_for_arc(vertebrae, s, e, dx_cobb, dx_ref_sign):
    idxs = np.arange(s, e + 1)
    apex_i = int(idxs[np.argmax(np.abs(dx_cobb[s:e + 1]))])

    upper_candidates = range(s, apex_i + 1)
    lower_candidates = range(apex_i, e + 1)

    best_cobb = -1.0
    best_pair = (s, e)

    for ui in upper_candidates:
        ta = vertebrae[ui]["ang"]
        for li in lower_candidates:
            c = float(angle_diff_deg(ta, vertebrae[li]["ang"]))
            if c > best_cobb:
                best_cobb = c
                best_pair = (ui, li)

    upper_idx, lower_idx = best_pair

    side_sign = arc_sign(dx_ref_sign, s, e)
    direction = "левосторонний" if side_sign > 0 else "правосторонний"

    return {
        "start_idx": int(s), "end_idx": int(e),
        "apex_idx": int(apex_i), "apex_label": vertebrae[apex_i]["l"],
        "direction": direction,
        "side_sign": int(side_sign),
        "cobb_deg": float(best_cobb),
        "upper_end": vertebrae[upper_idx]["l"],
        "lower_end": vertebrae[lower_idx]["l"],
        "upper_idx": int(upper_idx),
        "lower_idx": int(lower_idx),
    }

def arc_len(a):
    "Длина дуги в позвонках: end-start+1"
    return a["end_idx"] - a["start_idx"] + 1

def is_force(a, force_cobb, force_len):
    """"Форс-правило": если дуга очень большая по Cobb и не слишком короткая
    - считаем её структурной, даже если остальные критерии спорные."""
    return (a["cobb_deg"] >= force_cobb) and (arc_len(a) >= force_len)

def dedupe_by_range(arcs):
    "Функция удаляет дубликаты по (start_idx, end_idx)"
    seen = set()
    out = []
    for a in arcs:
        key = (a["start_idx"], a["end_idx"])
        if key not in seen:
            out.append(a)
            seen.add(key)
    return out

def split_arcs(
    arc_infos,
    # 1) Главная дуга всегда structural, если вообще есть "сколиоз"
    min_cobb_main=1.0,

    # 2) Вторая (для S-типа) должна быть не "шумом"
    min_cobb_second_abs=5.0,
    min_cobb_second_rel=0.35,

    min_len_struct=4,
    min_len_minor=3,
    min_apex_margin_struct=2,
    min_amp_rel_struct=0.25,

    min_cobb_minor=3.0,

    force_struct_cobb=40.0,
    force_struct_min_len=3
):
    """
    Разделяет все дуги на:
    - structural (важные, формируют тип C/S)
    - minor (дополнительные)
    """
    if not arc_infos:
        return [], []

    main = max(arc_infos, key=lambda a: a["cobb_deg"])

    def passes_geom(a):
        amp_rel = a.get("amp_rel")
        if amp_rel is None:
            return False

        apex_margin = int(a.get("apex_margin", 999))

        return (
            (arc_len(a) >= min_len_struct) and
            (apex_margin >= min_apex_margin_struct) and
            (amp_rel >= min_amp_rel_struct)
        )

    structural, minor = [], []

    main_is_force = (main["cobb_deg"] >= force_struct_cobb) and (arc_len(main) >= force_struct_min_len)

    if main_is_force or (main["cobb_deg"] >= min_cobb_main):
        structural.append(main)

    second_thr = max(min_cobb_second_abs, min_cobb_second_rel * float(main["cobb_deg"]))
    main_sign = int(main.get("side_sign", 0))

    for a in arc_infos:
        if a is main:
            continue

        length = arc_len(a)
        a_sign = int(a.get("side_sign", 0))

        if is_force(a, force_struct_cobb, force_struct_min_len):
            structural.append(a)
            continue

        is_opposite = (main_sign != 0) and (a_sign != 0) and (a_sign != main_sign)

        if is_opposite and (a["cobb_deg"] >= second_thr) and passes_geom(a):
            structural.append(a)
        else:
            if length >= min_len_minor and a["cobb_deg"] >= min_cobb_minor:
                minor.append(a)

    return dedupe_by_range(structural), dedupe_by_range(minor)

def chaklin_degree(cobb_deg):
    "Переводит Cobb в степень по Чаклину (0/I/II/III/IV)"
    a = float(cobb_deg)
    if a < 1.0:
        return "0 ст. (практически нет)"
    if a <= 10.:
        return "I ст."
    if a <= 25.:
        return "II ст."
    if a <= 50.:
        return "III ст."
    return "IV ст."

def scoliosis_type(structural_arcs):
    if len(structural_arcs) <= 1:
        return "C-сколиоз"

    signs = {int(a["side_sign"]) for a in structural_arcs if a.get("side_sign") is not None}

    if len(structural_arcs) >= 3 and len(signs) >= 2:
        return "Z-сколиоз"
    if len(signs) >= 2:
        return "S-сколиоз"
    return "C-сколиоз"

def build_conclusion(arc_infos, structural=None, minor=None):
    """
    Собирает текст заключения:
    - выбирает главную дугу
    - пишет тип, направление, Cobb, степень
    - выводит список структурных дуг
    - выводит список дополнительных дуг (minor)
    """
    if not arc_infos:
        return "Угловая деформация оси позвоночника не выявлена. Рентгенологических признаков сколиоза нет."
    if structural is None and minor is None:
        structural, minor = split_arcs(arc_infos)
    else:
        structural = structural or []
        minor = minor or []

    def fmt_arc(i, a, with_len=False):
        arc_degree = chaklin_degree(a["cobb_deg"])
        s = (f"  Дуга {i}: {a['direction']}, уровни {a['upper_end']}-{a['lower_end']}, "
             f"вершина ~ {a['apex_label']}, Cobb={a['cobb_deg']:.1f}°, степень={arc_degree}")
        if with_len:
            length = arc_len(a)
            s += f", длина={length} позв."
        return s + "."

    main = max(arc_infos, key=lambda x: x["cobb_deg"])

    scol_type = scoliosis_type(structural) if structural else "не определён (нет структурных дуг)"
    degree = chaklin_degree(main["cobb_deg"])

    lines = [
        f"1. Тип сколиоза - {scol_type}.",
        f"2. Направление основной дуги - {main['direction']} (по изображению).",
        f"3. Угол деформации (Cobb) основной дуги - {main['cobb_deg']:.1f}°.",
        f"4. Степень сколиоза по углу деформации - {degree}.",
        "",
        "Структурные дуги (учитываются для типа):"
    ]

    if structural:
        for i, a in enumerate(sorted(structural, key=lambda a: a["cobb_deg"], reverse=True), start=1):
            lines.append(fmt_arc(i, a))
    else:
        lines.append("  Нет (не прошли критерии структурности).")

    if minor:
        lines += ["", "Дополнительные дуги (не учитываются для типа):"]
        for i, a in enumerate(sorted(minor, key=lambda a: a["cobb_deg"], reverse=True), start=1):
            lines.append(fmt_arc(i, a, with_len=True))

    return "\n".join(lines)