scoliosis-inference/scoliosis_geometry.py

import numpy as np

def edge_angle_deg(p1, p2):
    "Находит угол наклона отрезка p1→p2 в градусах (от -180 до +180)."
    dx = float(p2[0] - p1[0])
    dy = float(p2[1] - p1[1])
    return np.degrees(np.arctan2(dy, dx))

def normalize_angle_minus90_90(a):
    "Приводит угол к диапазону [-90, 90)."
    return (a + 90.0) % 180.0 - 90.0

def angle_diff_deg(a, b):
    "Возвращает насколько отличаются две линии по углу, но только 'острый' вариант (0..90)"
    d = abs(a - b) % 180.0
    if d > 90.0:
        d = 180.0 - d
    return d

def order_points_ccw(pts):
    "Приводит 4 точки OBB к стабильному порядку по кругу"
    pts = np.asarray(pts, dtype=np.float32)
    c = pts.mean(axis=0)
    ang = np.arctan2(pts[:, 1] - c[1], pts[:, 0] - c[0])
    return pts[np.argsort(ang)]

def vertebra_endplates_angles(obb_pts):
    "Нахождит углов верхней и нижней замыкательных пластинок"
    p = order_points_ccw(obb_pts)

    edges = []
    for i in range(4):
        a = p[i]
        b = p[(i + 1) % 4]

        ang = normalize_angle_minus90_90(edge_angle_deg(a, b))
        mean_y = float((a[1] + b[1]) / 2.0)
        score = abs(ang)

        edges.append((score, mean_y, float(ang)))

    edges.sort(key=lambda t: t[0])
    e1, e2 = edges[0], edges[1]
    top, bottom = (e1, e2) if e1[1] <= e2[1] else (e2, e1)

    return top[2], bottom[2]

def fit_line_x_of_y(y, x, eps=1e-6):
    "Строит прямую вида x = a*y + b по точкам"
    if np.std(y) < eps:
        return 0.0, float(np.mean(x))
    a, b = np.polyfit(y, x, 1)
    return float(a), float(b)

def fit_spine_axis_dx(vertebrae):
    """Строит глобальную ось позвоночника (одна прямая по всем центрам).
    Возвращает dx: для каждого позвонка разницу показывая,
    насколько позвонок левее/правее глобальной оси"""
    centers = np.array([v["c"] for v in vertebrae], dtype=np.float32)
    xs, ys = centers[:, 0], centers[:, 1]

    a, b = fit_line_x_of_y(ys, xs)
    return xs - (a * ys + b)

def fit_spine_axis_dx_local(vertebrae, window=7):
    "Вычисляет локальные изгибы относительно соседей"
    centers = np.array([v["c"] for v in vertebrae], dtype=np.float32)
    xs, ys = centers[:, 0], centers[:, 1]
    n = len(xs)
    dx = np.zeros(n, dtype=np.float32)

    half = window // 2
    for i in range(n):
        s = max(0, i - half)
        e = min(n, i + half + 1)
        a, b = fit_line_x_of_y(ys[s:e], xs[s:e])
        dx[i] = xs[i] - (a * ys[i] + b)
    return dx

def smooth_1d(x, iters=3, alpha=0.25):
    "Простое сглаживание 1D-сигнала (dx)"
    x = x.astype(np.float32).copy()
    for _ in range(int(iters)):
        x[1:-1] = x[1:-1] + alpha * (x[:-2] - 2 * x[1:-1] + x[2:])
    return x

def fit_spine_axis_dx_endpoints(vertebrae):
    "Строит ось как линию через первый и последний позвонок (а не через регрессию)"
    centers = np.array([v["c"] for v in vertebrae], dtype=np.float32)
    xs, ys = centers[:, 0], centers[:, 1]

    x0, y0 = float(xs[0]), float(ys[0])
    x1, y1 = float(xs[-1]), float(ys[-1])

    if abs(y1 - y0) < 1e-6:
        return xs - float(np.mean(xs))

    t = (ys - y0) / (y1 - y0)
    x_line = x0 + t * (x1 - x0)
    return xs - x_line

def _smooth_chain(points, window=5):
    """сглаживает цепочку точек скользящим средним по окну window. 
    Делает окно нечётным, берёт среднее по соседям и возвращает сглаженные точки"""
    pts = np.asarray(points, dtype=np.float32)
    if len(pts) < 3 or window <= 2:
        return pts
    window = int(window)
    if window % 2 == 0:
        window += 1
    half = window // 2
    out = pts.copy()
    for i in range(half, len(pts) - half):
        out[i] = np.mean(pts[i - half:i + half + 1], axis=0)
    return out

def arc_centerline(vertebrae, info, smooth_window=5):
    """Строит центральную линию дуги"""
    s = info.get("upper_idx", info["start_idx"])
    e = info.get("lower_idx", info["end_idx"])
    if s > e:
        s, e = e, s
    centers = [vertebrae[k]["c"] for k in range(s, e + 1)]
    if not centers:
        return None
    centers = np.array(centers, dtype=np.float32)
    centers = centers[np.argsort(centers[:, 1])]
    if smooth_window and smooth_window > 2:
        centers = _smooth_chain(centers, window=smooth_window)
    return centers