roman_vinokurov/scoliosis-inference
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Initial scoliosis pipeline (without models)\

Browse Source
This commit is contained in:
Roman
2026-01-23 20:35:07 +05:00
commit 83dfbcc238
11 changed files with 1638 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
detect_vertebrae.py
+374
View File
@@ -0,0 +1,374 @@
import os
import cv2
import numpy as np
def detect_vertebrae(
image,
model_results,
labelmap,
debug_save_path=None,
enable_postprocessing=True,
interpolate_missing=False,
spine_sequence=None,
):
"""
Detect vertebrae with optional postprocessing.
Returns:
dict[label] = [x_center, y_center, width, height, angle, confidence]
"""
default_spine_sequence = [
"T1", "T2", "T3", "T4", "T5", "T6", "T7", "T8", "T9", "T10", "T11", "T12",
"L1", "L2", "L3", "L4", "L5",
]
if spine_sequence is None:
spine_sequence = default_spine_sequence
label_to_index = {label: idx for idx, label in enumerate(spine_sequence)}
best_boxes = {}
colors = [
(0, 0, 255), # Red
(0, 255, 0), # Green
(255, 0, 0), # Blue
(0, 255, 255), # Yellow
(255, 0, 255), # Magenta
(255, 255, 0), # Cyan
(0, 165, 255), # Orange
(255, 20, 147), # Pink
(147, 20, 255), # Violet
(0, 215, 255), # Gold
(255, 215, 0), # Turquoise
(255, 105, 180), # Hot pink
(0, 255, 127), # Spring green
(255, 69, 0), # Orange red
(72, 61, 139), # Dark slate blue
(47, 255, 173), # Aquamarine
(255, 140, 0), # Dark orange
]
if model_results:
result = model_results[0]
if hasattr(result, "obb") and result.obb is not None and len(result.obb) > 0:
obb_data = result.obb
for i in range(len(obb_data)):
try:
conf = float(obb_data.conf[i].cpu().numpy())
cls_id = int(obb_data.cls[i].cpu().numpy())
label = labelmap[cls_id] if cls_id < len(labelmap) else f"Unknown_{cls_id}"
if hasattr(obb_data, "xyxyxyxy"):
box_points = obb_data.xyxyxyxy[i].cpu().numpy()
elif hasattr(obb_data, "xyxyxyxyn"):
box_points = obb_data.xyxyxyxyn[i].cpu().numpy()
h, w = image.shape[:2]
box_points = box_points.reshape(4, 2)
box_points[:, 0] *= w
box_points[:, 1] *= h
box_points = box_points.flatten()
else:
box_points = result.obb.xyxyxyxy[i].cpu().numpy()
points = box_points.reshape(4, 2).astype(np.float32)
rect = cv2.minAreaRect(points)
center, size, angle = rect
x_center, y_center = center
width, height = size
if width < height:
angle += 90
width, height = height, width
result_box = np.array([x_center, y_center, width, height, angle, conf], dtype=np.float32)
if label not in best_boxes or conf > best_boxes[label][-1]:
best_boxes[label] = result_box
except Exception as exc:
print(f"Box parse error: {exc}")
continue
if enable_postprocessing and best_boxes:
boxes_list = [(label, box) for label, box in best_boxes.items()]
boxes_list.sort(key=lambda x: x[1][-1], reverse=True)
kept_boxes = []
for label, box in boxes_list:
x_center, y_center, width, height, _, _ = box
is_duplicate = False
for kept_label, kept_box in kept_boxes:
xk, yk, wk, hk, _, _ = kept_box
dist = np.sqrt((x_center - xk) ** 2 + (y_center - yk) ** 2)
threshold = 0.3 * min(height, hk)
if dist < threshold:
is_duplicate = True
print(f"Duplicate removed: {label} overlaps with {kept_label}")
break
if not is_duplicate:
kept_boxes.append((label, box))
kept_boxes.sort(key=lambda x: x[1][1])
if kept_boxes:
valid_boxes = []
last_index = -1
last_y = -1
y_gaps = []
for label, box in kept_boxes:
x_center, y_center, width, height, angle, conf = box
if label not in label_to_index:
print(f"Unknown vertebra skipped: {label}")
continue
current_index = label_to_index[label]
if last_index == -1:
valid_boxes.append((label, box))
last_index = current_index
last_y = y_center
continue
expected_index = last_index + 1
if current_index == expected_index:
valid_boxes.append((label, box))
y_gaps.append(y_center - last_y)
last_index = current_index
last_y = y_center
continue
if current_index > expected_index:
valid_boxes.append((label, box))
last_index = current_index
last_y = y_center
continue
avg_gap = np.median(y_gaps) if y_gaps else 50.0
expected_y = last_y + avg_gap
y_diff = abs(y_center - expected_y)
if y_diff < avg_gap * 0.6 and expected_index < len(spine_sequence):
new_label = spine_sequence[expected_index]
new_box = box.copy()
new_box[5] = -1.0
valid_boxes.append((new_label, new_box))
y_gaps.append(y_center - last_y)
print(f"Order corrected: {label} -> {new_label} (conf=-1)")
last_index = expected_index
last_y = y_center
continue
print(f"Out-of-order skipped: {label} after {valid_boxes[-1][0]}")
if interpolate_missing and len(valid_boxes) > 1:
valid_boxes.sort(key=lambda x: x[1][1])
gaps = []
for i in range(1, len(valid_boxes)):
prev_y = valid_boxes[i - 1][1][1]
curr_y = valid_boxes[i][1][1]
gaps.append(curr_y - prev_y)
avg_gap = np.median(gaps) if gaps else 0
max_allowed_gap = avg_gap * 1.8 if avg_gap > 0 else float("inf")
reliable_angles = []
reliable_ratios = []
for _, box in valid_boxes:
_, _, width, height, angle, conf = box
if conf > 0:
norm_angle = angle % 180
if norm_angle > 90:
norm_angle -= 180
reliable_angles.append(norm_angle)
aspect_ratio = width / max(height, 1)
reliable_ratios.append(aspect_ratio)
median_angle = np.median(reliable_angles) if reliable_angles else 0.0
median_ratio = np.median(reliable_ratios) if reliable_ratios else 2.0
new_boxes = []
for i in range(len(valid_boxes)):
current_label, current_box = valid_boxes[i]
current_idx = label_to_index[current_label]
new_boxes.append((current_label, current_box))
if i < len(valid_boxes) - 1:
next_label, next_box = valid_boxes[i + 1]
next_idx = label_to_index[next_label]
y_gap = next_box[1] - current_box[1]
index_gap = next_idx - current_idx
if index_gap > 1 and y_gap > max_allowed_gap * 0.7:
num_missing = index_gap - 1
print(f"Missing between {current_label} and {next_label}: {num_missing}")
x1, y1, w1, h1, ang1, _ = current_box
x2, y2, w2, h2, ang2, _ = next_box
for k in range(1, num_missing + 1):
missing_idx = current_idx + k
if missing_idx >= len(spine_sequence):
continue
missing_label = spine_sequence[missing_idx]
fraction = k / index_gap
x_center = (x1 + x2) / 2
y_center = y1 + fraction * (y2 - y1)
avg_width = (w1 + w2) / 2
avg_height = (h1 + h2) / 2
if median_ratio > 1:
width = avg_width
height = width / median_ratio
else:
height = avg_height
width = height * median_ratio
norm_ang1 = ang1 % 180
norm_ang2 = ang2 % 180
if abs(norm_ang1 - norm_ang2) > 90:
if norm_ang1 > norm_ang2:
norm_ang1 -= 180
else:
norm_ang2 -= 180
angle = norm_ang1 + fraction * (norm_ang2 - norm_ang1)
angle = (angle + 90) % 180 - 90
if reliable_angles:
angle = 0.7 * angle + 0.3 * median_angle
if height > width:
width, height = height, width
angle = (angle + 90) % 180
conf = -0.5
interpolated_box = np.array(
[x_center, y_center, width, height, angle, conf],
dtype=np.float32,
)
print(
"Interpolated: "
f"{missing_label} x={x_center:.1f} y={y_center:.1f} "
f"w={width:.1f} h={height:.1f} angle={angle:.2f}"
)
new_boxes.append((missing_label, interpolated_box))
new_boxes.sort(key=lambda x: x[1][1])
valid_boxes = new_boxes
best_boxes = {label: box for label, box in valid_boxes}
if debug_save_path:
if len(image.shape) == 2:
vis_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
elif image.shape[2] == 3:
vis_image = image.copy()
else:
vis_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
for idx, (label, box) in enumerate(best_boxes.items()):
try:
x_center, y_center, width, height, angle, conf = box
color = colors[idx % len(colors)]
is_corrected = conf < 0
thickness = 4 if is_corrected else 3
alpha = 0.35 if is_corrected else 0.2
rect = ((x_center, y_center), (width, height), angle)
box_points = cv2.boxPoints(rect)
box_points = np.int64(box_points)
contour_color = (0, 0, 255) if is_corrected else color
cv2.drawContours(vis_image, [box_points], 0, contour_color, thickness)
overlay = vis_image.copy()
fill_color = (0, 0, 255) if is_corrected else color
cv2.fillPoly(overlay, [box_points], fill_color)
cv2.addWeighted(overlay, alpha, vis_image, 1 - alpha, 0, vis_image)
if is_corrected:
text = f"{label} CV({conf:.1f})"
else:
text = f"{label} ({conf:.3f})"
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.6
thickness = 2
(text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
text_x = int(x_center - text_width / 2)
text_y = int(y_center + text_height / 2)
padding = 5
bg_rect = [
(text_x - padding, text_y - text_height - padding),
(text_x + text_width + padding, text_y + padding),
]
text_overlay = vis_image.copy()
bg_color = (0, 0, 100) if is_corrected else (0, 0, 0)
cv2.rectangle(text_overlay, bg_rect[0], bg_rect[1], bg_color, -1)
text_alpha = 0.7
cv2.addWeighted(text_overlay, text_alpha, vis_image, 1 - text_alpha, 0, vis_image)
text_color = (255, 255, 255) if not is_corrected else (255, 200, 200)
cv2.putText(
vis_image,
text,
(text_x, text_y),
font,
font_scale,
text_color,
thickness,
)
center_color = (0, 0, 255) if is_corrected else (255, 255, 255)
cv2.circle(vis_image, (int(x_center), int(y_center)), 5, center_color, -1)
cv2.circle(vis_image, (int(x_center), int(y_center)), 2, (0, 0, 0), -1)
except Exception as exc:
print(f"Draw error for {label}: {exc}")
continue
dir_path = os.path.dirname(debug_save_path)
if dir_path:
os.makedirs(dir_path, exist_ok=True)
try:
if len(vis_image.shape) == 2:
vis_image = cv2.cvtColor(vis_image, cv2.COLOR_GRAY2BGR)
cv2.imwrite(debug_save_path, vis_image, [cv2.IMWRITE_JPEG_QUALITY, 95])
print(f"Debug saved: {debug_save_path}")
print(f"Detections: {len(best_boxes)}")
corrected_count = sum(1 for box in best_boxes.values() if box[5] < 0)
if corrected_count > 0:
print(f"Corrected order: {corrected_count}")
except Exception as exc:
print(f"Debug save error: {exc}")
return {label: box.astype(float).tolist() for label, box in best_boxes.items()}