Roman
275 lines
10 KiB
Python
275 lines
10 KiB
Python
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)
|