Initial scoliosis pipeline (without models)\
This commit is contained in:
Roman
+274
@@ -0,0 +1,274 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from classification import classify_projection, classify_stitched
|
||||
from config import load_config
|
||||
from detect_vertebrae import detect_vertebrae
|
||||
from io_dicom import read_dicom, save_dicom_sr, save_secondary_capture
|
||||
from model_loader import (
|
||||
get_classification_model,
|
||||
get_detection_model,
|
||||
get_projection_model,
|
||||
)
|
||||
from report_builder import build_scoliosis_description
|
||||
from scoliosis_arcs import (
|
||||
build_conclusion,
|
||||
cobb_for_arc,
|
||||
detect_arcs_from_dx,
|
||||
merge_arcs,
|
||||
split_arcs,
|
||||
scoliosis_type,
|
||||
)
|
||||
from scoliosis_geometry import (
|
||||
arc_centerline,
|
||||
fit_spine_axis_dx,
|
||||
fit_spine_axis_dx_endpoints,
|
||||
fit_spine_axis_dx_local,
|
||||
smooth_1d,
|
||||
)
|
||||
from vis import draw_arc_contours
|
||||
|
||||
CLASS_NAMES = {
|
||||
0: "L1", 1: "L2", 2: "L3", 3: "L4", 4: "L5",
|
||||
5: "T1", 6: "T2", 7: "T3", 8: "T4", 9: "T5",
|
||||
10: "T6", 11: "T7", 12: "T8", 13: "T9", 14: "T10",
|
||||
15: "T11", 16: "T12",
|
||||
}
|
||||
|
||||
CONFIG = load_config()
|
||||
POSTPROCESSING = CONFIG.postprocessing
|
||||
INTERPOLATE_MISSING = CONFIG.interpolate_missing
|
||||
|
||||
def build_labelmap(class_names):
|
||||
"""Делает список меток по индексам (0..max),
|
||||
чтобы по cls_id получить имя позвонка."""
|
||||
max_id = max(class_names)
|
||||
return [class_names.get(i, f"Unknown_{i}") for i in range(max_id + 1)]
|
||||
|
||||
def vertebrae_from_detect_result(detect_result):
|
||||
"""Преобразует результат детекции в список позвонков:
|
||||
центр, коробка, угол, уверенность; сортирует по Y."""
|
||||
vertebrae = []
|
||||
for label, box in detect_result.items():
|
||||
x_center, y_center, width, height, angle, conf = box
|
||||
rect = (
|
||||
(float(x_center), float(y_center)),
|
||||
(float(width), float(height)),
|
||||
float(angle),
|
||||
)
|
||||
pts = cv2.boxPoints(rect).astype(np.float32)
|
||||
center = np.array([x_center, y_center], dtype=np.float32)
|
||||
vertebrae.append({
|
||||
"id": label,
|
||||
"l": label,
|
||||
"c": center,
|
||||
"b": pts,
|
||||
"ang": float(angle),
|
||||
"conf": float(conf),
|
||||
})
|
||||
vertebrae.sort(key=lambda x: x["c"][1])
|
||||
return vertebrae
|
||||
|
||||
def run_pipeline(img_bgr, vertebrae):
|
||||
"""Основная логика анализа: строит оси/сигналы, находит дуги, считает Cobb,
|
||||
делит дуги на структурные/минорные, формирует заключение и итоговые метрики."""
|
||||
if img_bgr is None:
|
||||
raise ValueError("img_bgr is None")
|
||||
|
||||
h, w = img_bgr.shape[:2]
|
||||
|
||||
dx_global = fit_spine_axis_dx(vertebrae)
|
||||
dx_seg = smooth_1d(dx_global, iters=3, alpha=0.25)
|
||||
|
||||
dx_local = fit_spine_axis_dx_local(vertebrae, window=7)
|
||||
dx_cobb = smooth_1d(dx_local, iters=3, alpha=0.25)
|
||||
|
||||
dx_ref = smooth_1d(fit_spine_axis_dx_endpoints(vertebrae), iters=3, alpha=0.25)
|
||||
|
||||
arcs = detect_arcs_from_dx(
|
||||
dx_cobb,
|
||||
eps_px=CONFIG.detect_eps_px,
|
||||
min_len=CONFIG.detect_min_len,
|
||||
zero_gap=CONFIG.detect_zero_gap,
|
||||
edge_zero_extend=CONFIG.detect_edge_zero_extend,
|
||||
)
|
||||
arcs = merge_arcs(arcs, dx_ref, max_gap=CONFIG.merge_max_gap)
|
||||
|
||||
global_amp = float(np.max(np.abs(dx_seg)) + 1e-6)
|
||||
|
||||
arc_infos = []
|
||||
for (s, e) in arcs:
|
||||
info = cobb_for_arc(vertebrae, s, e, dx_cobb, dx_ref)
|
||||
info["amp_rel"] = float(np.max(np.abs(dx_seg[s:e + 1])) / global_amp)
|
||||
info["apex_margin"] = int(min(info["apex_idx"] - s, e - info["apex_idx"]))
|
||||
info["centerline_pts"] = arc_centerline(vertebrae, info, smooth_window=5)
|
||||
arc_infos.append(info)
|
||||
|
||||
n = len(vertebrae)
|
||||
min_len_struct = CONFIG.min_len_struct_small if n <= 7 else CONFIG.min_len_struct_large
|
||||
min_apex_margin_struct = CONFIG.min_apex_margin_small if n <= 7 else CONFIG.min_apex_margin_large
|
||||
|
||||
structural, minor = split_arcs(
|
||||
arc_infos,
|
||||
min_cobb_main=CONFIG.min_cobb_main,
|
||||
min_cobb_second_abs=CONFIG.min_cobb_second_abs,
|
||||
min_cobb_second_rel=CONFIG.min_cobb_second_rel,
|
||||
min_len_struct=min_len_struct,
|
||||
min_apex_margin_struct=min_apex_margin_struct,
|
||||
)
|
||||
|
||||
conclusion = build_conclusion(arc_infos, structural=structural, minor=minor)
|
||||
|
||||
main = max(arc_infos, key=lambda x: x["cobb_deg"]) if arc_infos else None
|
||||
degree_class = 0
|
||||
prob_scoliosis = 0.0
|
||||
if main:
|
||||
cd = main["cobb_deg"]
|
||||
if cd >= 50:
|
||||
degree_class = 4
|
||||
elif cd >= 26:
|
||||
degree_class = 3
|
||||
elif cd >= 11:
|
||||
degree_class = 2
|
||||
elif cd >= 1:
|
||||
degree_class = 1
|
||||
prob_scoliosis = min(1.0, cd / 50.0)
|
||||
|
||||
scol_type = "C-сколиоз"
|
||||
if structural:
|
||||
scol_type = scoliosis_type(structural)
|
||||
else:
|
||||
scol_type = "нет сколиоза"
|
||||
|
||||
return {
|
||||
"img_bgr": img_bgr,
|
||||
"h": h, "w": w,
|
||||
"vertebrae": vertebrae,
|
||||
"arc_infos": arc_infos,
|
||||
"structural": structural,
|
||||
"minor": minor,
|
||||
"conclusion": conclusion,
|
||||
"main_arc": main,
|
||||
"degree_class": degree_class,
|
||||
"scoliosis_type": scol_type,
|
||||
"prob_scoliosis": prob_scoliosis,
|
||||
}
|
||||
|
||||
def _split_if_stitched(img_bgr, cls_model):
|
||||
"""Проверяет, “сшитый” ли снимок.
|
||||
Если да — делит на левую/правую половины,
|
||||
иначе возвращает исходный."""
|
||||
pred, conf = classify_stitched(cls_model, img_bgr)
|
||||
if pred != 1: # 0 = single
|
||||
return [img_bgr], {"classification": pred, "classification_conf": conf}
|
||||
h, w = img_bgr.shape[:2]
|
||||
mid = w // 2
|
||||
left = img_bgr[:, :mid].copy()
|
||||
right = img_bgr[:, mid:].copy()
|
||||
return [left, right], {"classification": pred, "classification_conf": conf}
|
||||
|
||||
def _select_frontal(images, proj_model):
|
||||
"""Из списка изображений выбирает фронтальную проекцию;
|
||||
если нет — берёт с максимальной уверенностью."""
|
||||
best_img = images[0]
|
||||
best_meta = {"projection": None, "projection_conf": None}
|
||||
for img in images:
|
||||
pred, conf = classify_projection(proj_model, img)
|
||||
if pred == 1: # 1 = frontal
|
||||
return img, {"projection": pred, "projection_conf": conf}
|
||||
if best_meta["projection_conf"] is None or conf > best_meta["projection_conf"]:
|
||||
best_img, best_meta = img, {"projection": pred, "projection_conf": conf}
|
||||
return best_img, best_meta
|
||||
|
||||
def run_scoliosis_pipeline(infer_dicom_path, model=None, model_path=None, out_dir=None):
|
||||
"""Полный пайплайн: читает DICOM, выбирает проекцию, детектит позвонки,
|
||||
запускает run_pipeline, сохраняет SC и SR DICOM, возвращает пути и результат."""
|
||||
model_path = model_path or CONFIG.models.detection
|
||||
base_results_dir = out_dir or CONFIG.results_dir
|
||||
image_stem = os.path.splitext(os.path.basename(infer_dicom_path))[0]
|
||||
run_dir = os.path.join(base_results_dir, f"{image_stem}_result")
|
||||
os.makedirs(run_dir, exist_ok=True)
|
||||
|
||||
det_model = model or get_detection_model(model_path=model_path)
|
||||
cls_model = get_classification_model()
|
||||
proj_model = get_projection_model()
|
||||
|
||||
img_bgr_full = read_dicom(infer_dicom_path)
|
||||
split_imgs, cls_meta = _split_if_stitched(img_bgr_full, cls_model)
|
||||
img_bgr, proj_meta = _select_frontal(split_imgs, proj_model)
|
||||
print(f"classification pred={cls_meta['classification']} conf={cls_meta['classification_conf']:.3f}; "
|
||||
f"projection pred={proj_meta['projection']} conf={proj_meta['projection_conf']:.3f}")
|
||||
|
||||
yolo_results = det_model(img_bgr, verbose=False)
|
||||
|
||||
labelmap = build_labelmap(CLASS_NAMES)
|
||||
detect_result = detect_vertebrae(
|
||||
img_bgr,
|
||||
yolo_results,
|
||||
labelmap,
|
||||
enable_postprocessing=POSTPROCESSING,
|
||||
interpolate_missing=INTERPOLATE_MISSING,
|
||||
)
|
||||
|
||||
if not detect_result:
|
||||
return {
|
||||
"ok": False,
|
||||
"reason": "no_detections",
|
||||
}
|
||||
|
||||
vertebrae = vertebrae_from_detect_result(detect_result)
|
||||
result = run_pipeline(img_bgr, vertebrae)
|
||||
|
||||
sc_img = draw_arc_contours(img_bgr, vertebrae, result["structural"])
|
||||
sc_out_path = os.path.join(run_dir, f"sc_{image_stem}.dcm")
|
||||
overlay_txt = f"prob={result.get('prob_scoliosis', 0.0):.2f} deg={result.get('degree_class', 0)} type={result.get('scoliosis_type', '')}"
|
||||
save_secondary_capture(
|
||||
infer_dicom_path,
|
||||
sc_img,
|
||||
sc_out_path,
|
||||
series_description="RG-SCOLIOSIS-CONTOUR",
|
||||
overlay_text=overlay_txt,
|
||||
)
|
||||
|
||||
sr_desc = build_scoliosis_description(result["arc_infos"])
|
||||
sr_out_path = os.path.join(run_dir, f"sr_{image_stem}.dcm")
|
||||
main_cobb = result["main_arc"]["cobb_deg"] if result.get("main_arc") else 0.0
|
||||
save_dicom_sr(
|
||||
infer_dicom_path,
|
||||
sr_desc,
|
||||
result["conclusion"],
|
||||
sr_out_path,
|
||||
series_description="RG-SCOLIOSIS-SR",
|
||||
extra_struct={
|
||||
"degree_class": result.get("degree_class", 0),
|
||||
"scoliosis_type": result.get("scoliosis_type", ""),
|
||||
"main_cobb": main_cobb,
|
||||
"prob_scoliosis": result.get("prob_scoliosis", 0.0),
|
||||
},
|
||||
)
|
||||
|
||||
return {
|
||||
"ok": True,
|
||||
"sr_dicom_path": sr_out_path,
|
||||
"sc_dicom_path": sc_out_path,
|
||||
"conclusion": result["conclusion"],
|
||||
"arc_infos": result["arc_infos"],
|
||||
"out_dir": run_dir,
|
||||
}
|
||||
|
||||
if __name__ == "__main__":
|
||||
MODEL_PATH = CONFIG.models.detection
|
||||
INFER_IMG_PATH = os.environ.get("INFER_IMG_PATH", r"C:\Users\Роман Владимирович\Desktop\XR_SCOLIOS\XR_SCOLIOS\1.2.643.5.1.13.13.12.2.77.8252.00090213030609010011090905030205\1.2.643.5.1.13.13.12.2.77.8252.09051310090608150606061508100113\1.2.643.5.1.13.13.12.2.77.8252.03061304011101150513090811030403.dcm")
|
||||
OUT_DIR = CONFIG.results_dir
|
||||
|
||||
if not INFER_IMG_PATH:
|
||||
sys.exit(0)
|
||||
|
||||
model = get_detection_model(model_path=MODEL_PATH)
|
||||
out = run_scoliosis_pipeline(INFER_IMG_PATH, model=model, out_dir=OUT_DIR)
|
||||
|
||||
if not out.get("ok"):
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user