Aleksandr Mochalov
485 lines
18 KiB
Python
485 lines
18 KiB
Python
import warnings
|
|
from typing import NamedTuple, Optional
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.ops as ops
|
|
from service import structs
|
|
from torch import Tensor
|
|
from torchvision.transforms import v2 as T
|
|
|
|
warnings.filterwarnings('ignore', category=UserWarning)
|
|
device = torch.device('cuda')
|
|
|
|
models_root = "service/models/wrist/"
|
|
model_frac = torch.load(f'{models_root}/frac_model.pth',
|
|
weights_only=False).to(device)
|
|
model_bone = torch.load(f'{models_root}/bone_model.pth',
|
|
weights_only=False).to(device)
|
|
model_lr = torch.load(f'{models_root}/lr_model.pth',
|
|
weights_only=False).to(device)
|
|
model_move = torch.load(f'{models_root}/move_model.pth',
|
|
weights_only=False).to(device)
|
|
|
|
model_frac.eval()
|
|
model_bone.eval()
|
|
model_lr.eval()
|
|
model_move.eval()
|
|
|
|
|
|
class ProjectionSegments(NamedTuple):
|
|
parts_boxes: Tensor
|
|
parts_labels: list[str]
|
|
bones_boxes: Tensor
|
|
bones_labels: list[str]
|
|
|
|
|
|
class Fractures(NamedTuple):
|
|
boxes: Tensor
|
|
scores: Tensor
|
|
labels: list[str]
|
|
bones: list[str]
|
|
|
|
|
|
class CLAHETransform:
|
|
def __init__(self, clipLimit=2.0, tileGridSize=(8, 8)):
|
|
self.clipLimit = clipLimit
|
|
self.tileGridSize = tileGridSize
|
|
self.clahe = cv2.createCLAHE(self.clipLimit, self.tileGridSize)
|
|
|
|
def __call__(self, img):
|
|
img_np = img.cpu().numpy()
|
|
if img_np.ndim == 3 and img_np.shape[0] == 1:
|
|
img_np = img_np[0]
|
|
cl_img = self.clahe.apply(img_np)
|
|
return torch.from_numpy(cl_img).unsqueeze(0).to(img.device)
|
|
|
|
|
|
transform_bone = T.Compose([
|
|
T.Resize((256, 256)),
|
|
CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)),
|
|
T.ToDtype(torch.float, scale=True),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.3354] * 3, std=[0.2000] * 3)
|
|
])
|
|
|
|
transform_frac = T.Compose([
|
|
T.Resize((512, 512)),
|
|
CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)),
|
|
T.ToDtype(torch.float, scale=True),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.21549856662750244], std=[0.24515700340270996])
|
|
])
|
|
|
|
transform_lr = T.Compose([
|
|
T.Resize((256, 256)),
|
|
T.ToDtype(torch.float, scale=True),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.9278] * 3, std=[0.2089] * 3)
|
|
])
|
|
|
|
transform_move = T.Compose([
|
|
T.Resize((224, 224)),
|
|
CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)),
|
|
T.ToDtype(torch.float, scale=True),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.3354] * 3, std=[0.2000] * 3)
|
|
])
|
|
|
|
bone_translator = {
|
|
'1': 'локтевой кости',
|
|
'2': 'лучевой кости',
|
|
'3': 'лучевой или локтевой кости',
|
|
'4': 'костей запястья'
|
|
}
|
|
bone_parts_translator = {
|
|
'head': 'головки',
|
|
'styloid': 'шиловидного отростка',
|
|
'epiphysis': 'эпифиза',
|
|
'metaphysis': 'метафиза',
|
|
'diaphysis': 'диафиза',
|
|
'hand': ''
|
|
}
|
|
|
|
bone_translator_clear = {
|
|
'1': 'Локтевая кость',
|
|
'2': 'Лучевая кость',
|
|
'3': 'Лучевая или локтевая кость',
|
|
'4': 'Кости запястья'
|
|
}
|
|
|
|
bone_ratios = {
|
|
'1': {'styloid_1': 0.02, 'head_1': 0.02, 'epiphysis_1': 0.05,
|
|
'metaphysis_1': 0.03, 'diaphysis_1': 0.8},
|
|
'2': {'styloid_2': 0.04, 'epiphysis_2': 0.13, 'metaphysis_2': 0.03,
|
|
'diaphysis_2': 0.8},
|
|
'3': {'styloid_3': 0.04, 'epiphysis_3': 0.13, 'metaphysis_3': 0.03,
|
|
'diaphysis_3': 0.8}
|
|
}
|
|
|
|
|
|
def _convert_bboxes(bboxes: Tensor, orig_w: int, orig_h: int,
|
|
transformed_w=256, transformed_h=256) -> Tensor:
|
|
"""Масштабирует координаты боксов обратно к
|
|
размерам исходного изображения."""
|
|
bboxes[:, [0, 2]] *= orig_w / transformed_w
|
|
bboxes[:, [1, 3]] *= orig_h / transformed_h
|
|
return bboxes
|
|
|
|
|
|
def _get_bone_parts(box: Tensor, bone_type: str) -> dict[str: Tensor]:
|
|
"""Разбивает бокc кости на сегменты согласно заданным пропорциям."""
|
|
if bone_type == '4':
|
|
return {'hand_4': box}
|
|
xmin, ymin, xmax = box[:3]
|
|
bone_h = (xmax - xmin) * (9 if bone_type == '1' else 5)
|
|
parts, current_y = {}, ymin
|
|
for part, ratio in bone_ratios[bone_type].items():
|
|
part_h = ratio * bone_h
|
|
parts[part] = (xmin, current_y, xmax, current_y + part_h)
|
|
current_y = current_y + part_h
|
|
return parts
|
|
|
|
|
|
def _check_is_right(image: Tensor) -> bool:
|
|
"""Определяет сторону (латеральность) по изображению."""
|
|
image = transform_lr(image).unsqueeze(0).to(device)
|
|
with torch.no_grad():
|
|
outputs = model_lr(image)[0]
|
|
return bool(torch.argmax(outputs).item())
|
|
|
|
|
|
def _get_projection_bones(image: Tensor) -> Optional[ProjectionSegments]:
|
|
"""Получает боксы и сегменты кости для проекции."""
|
|
original = image.clone()
|
|
image = transform_bone(image).unsqueeze(0).to(device)
|
|
with torch.no_grad():
|
|
outputs = model_bone(image)[0]
|
|
scores = outputs['scores']
|
|
|
|
valid = scores >= 0.15
|
|
bones_boxes = outputs['boxes'][valid].cpu()
|
|
bone_labels = [str(int(i.item())) for i in outputs['labels'][valid].cpu()]
|
|
filtered_scores = scores[valid]
|
|
|
|
if not any(filtered_scores):
|
|
raise structs.ImagesError("Снимок иной анатомической области или низкого диагностического качества")
|
|
|
|
best = {}
|
|
for score, box, label in zip(filtered_scores, bones_boxes, bone_labels):
|
|
if label not in best or score.item() > best[label][0]:
|
|
best[label] = (score.item(), box)
|
|
# 4 - ладонь
|
|
if '3' in best:
|
|
best['4'] = best['3']
|
|
best.pop('3')
|
|
bones_boxes = torch.stack([b for _, b in best.values()])
|
|
bone_labels = list(best.keys())
|
|
conv_bone_boxes = _convert_bboxes(bones_boxes, original.shape[2],
|
|
original.shape[1])
|
|
|
|
parts_boxes, parts_labels = [], []
|
|
for box, bone in zip(conv_bone_boxes, bone_labels):
|
|
segments = _get_bone_parts(box.cpu().numpy(), bone)
|
|
for part, coords in segments.items():
|
|
parts_boxes.append(torch.tensor(coords))
|
|
parts_labels.append(part)
|
|
parts_boxes = torch.stack(parts_boxes)
|
|
|
|
return ProjectionSegments(parts_boxes, parts_labels, conv_bone_boxes,
|
|
bone_labels)
|
|
|
|
|
|
def _get_fractions(direct_img: Tensor,
|
|
score_threshold=0.35) -> tuple[Tensor, Tensor]:
|
|
image = transform_frac(direct_img).unsqueeze(0).to(device)
|
|
with torch.no_grad():
|
|
outputs = model_frac(image)[0]
|
|
scores = outputs['scores'].cpu()
|
|
valid = scores >= score_threshold
|
|
boxes = outputs['boxes'][valid].cpu()
|
|
scores = scores[valid]
|
|
keep = ops.nms(boxes, scores, 0.2)
|
|
boxes = boxes[keep]
|
|
scores = scores[keep]
|
|
converted_bboxes = _convert_bboxes(boxes, direct_img.shape[2],
|
|
direct_img.shape[1],
|
|
transformed_w=512,
|
|
transformed_h=512)
|
|
return converted_bboxes, scores
|
|
|
|
|
|
def _assign_fracs_to_bones(frac_boxes: Tensor,
|
|
bone_boxes: Tensor, bone_labels: list[str]) \
|
|
-> dict[str: tuple[Tensor, list[str]]]:
|
|
"""Относит каждый бокс ровно к одной кости"""
|
|
assignments = {}
|
|
new_boxes = []
|
|
i = 1
|
|
if len(bone_boxes) > 0:
|
|
for frac_box in frac_boxes:
|
|
inter_xmin = torch.max(frac_box[0], bone_boxes[:, 0])
|
|
inter_ymin = torch.max(frac_box[1], bone_boxes[:, 1])
|
|
inter_xmax = torch.min(frac_box[2], bone_boxes[:, 2])
|
|
inter_ymax = torch.min(frac_box[3], bone_boxes[:, 3])
|
|
|
|
x_len = (inter_xmax - inter_xmin).clamp(min=0)
|
|
y_len = (inter_ymax - inter_ymin).clamp(min=0)
|
|
inter_area = x_len * y_len
|
|
|
|
best_idx = torch.argmax(inter_area)
|
|
if inter_area[best_idx].item() > 0:
|
|
bone = bone_labels[best_idx]
|
|
if bone in assignments:
|
|
assignments[bone] = (
|
|
torch.cat([assignments[bone][0], frac_box.unsqueeze(0)]),
|
|
assignments[bone][1] + [f'Находка {i}'])
|
|
else:
|
|
assignments[bone] = (frac_box.unsqueeze(0), [f'Находка {i}'])
|
|
new_boxes.append(frac_box)
|
|
|
|
i += 1
|
|
boxes_tensor = torch.stack(new_boxes) if new_boxes else torch.tensor([])
|
|
else:
|
|
boxes_tensor = torch.tensor([])
|
|
return assignments, boxes_tensor
|
|
|
|
|
|
def _assign_fracs_to_parts(frac_boxes: Tensor, frac_labels: list[str],
|
|
part_boxes: Tensor,
|
|
part_labels: np.ndarray[str]) -> list[str]:
|
|
"""Для каждого перелома ищет 2 части кости с наибольшим пересечением"""
|
|
result = []
|
|
|
|
for i in range(frac_boxes.shape[0]):
|
|
frac_box = frac_boxes[i]
|
|
intersections = []
|
|
for j in range(part_boxes.shape[0]):
|
|
part_box = part_boxes[j]
|
|
inter_xmin = max(frac_box[0].item(), part_box[0].item())
|
|
inter_ymin = max(frac_box[1].item(), part_box[1].item())
|
|
inter_xmax = min(frac_box[2].item(), part_box[2].item())
|
|
inter_ymax = min(frac_box[3].item(), part_box[3].item())
|
|
|
|
# Если пересечение существует, вычисляем площадь пересечения
|
|
if inter_xmax > inter_xmin and inter_ymax > inter_ymin:
|
|
area = (inter_xmax - inter_xmin) * (inter_ymax - inter_ymin)
|
|
intersections.append((j, area))
|
|
|
|
# Сортируем найденные пересечения по площади в
|
|
# порядке убывания и выбираем топ-2
|
|
intersections.sort(key=lambda x: x[1], reverse=True)
|
|
for j, _ in intersections[:2]:
|
|
result.append(f'{part_labels[j]}_{frac_labels[i]}')
|
|
|
|
return result
|
|
|
|
|
|
def _format_single_fracture(label: str) -> tuple[str, str]:
|
|
"""
|
|
Формирует сообщение для одного найденного перелома.
|
|
"""
|
|
name, part = label.split('_')[:2]
|
|
bone = bone_translator[part]
|
|
msg = f'выявлен признак перелома {bone_parts_translator[name]} {bone}.'
|
|
return ('', msg) if part == '3' else (msg, '')
|
|
|
|
|
|
def _format_multiple_fractures(labels: list[str]) -> tuple[str, str, int, int]:
|
|
"""
|
|
Группирует найденные переломы по идентификатору и формирует сообщения
|
|
для фронтальной и боковой проекций.
|
|
"""
|
|
finds = {}
|
|
for label in labels:
|
|
name, part, n = label.split('_')
|
|
bone = bone_translator[part]
|
|
finds.setdefault(n, []).append((bone_parts_translator[name], bone))
|
|
|
|
front_msgs, side_msgs = [], []
|
|
front_count = side_count = 0
|
|
for n in sorted(finds.keys()):
|
|
parts = [x[0] for x in finds[n]]
|
|
if 'метафиза' in parts and 'эпифиза' in parts:
|
|
parts = ['метаэпифиза']
|
|
find_text = f' {n}: перелом ' + ', '.join(parts) + f' {finds[n][0][1]}.'
|
|
if finds[n][0][1] != 'лучевой или локтевой кости':
|
|
front_msgs.append(find_text)
|
|
front_count += 1
|
|
else:
|
|
side_msgs.append(find_text)
|
|
side_count += 1
|
|
|
|
front_msg = (f'выявлены признаки {front_count} ' +
|
|
('переломов' if front_count > 1 else 'перелома') +
|
|
'. ' + ''.join(front_msgs)) if front_count else ''
|
|
side_msg = (f'выявлены признаки {side_count} ' +
|
|
('переломов' if side_count > 1 else 'перелома') +
|
|
'. ' + ''.join(side_msgs)) if side_count else ''
|
|
return front_msg, side_msg, front_count, side_count
|
|
|
|
|
|
def _make_reports(labels: list[str], is_r: bool) -> str:
|
|
"""
|
|
Формирует текстовый отчёт по найденным переломам.
|
|
"""
|
|
|
|
labels = list(set(labels))
|
|
num_frac = len(labels)
|
|
lr = 'правого' if is_r else 'левого'
|
|
|
|
# Отсутствие переломов
|
|
if num_frac == 0:
|
|
return (
|
|
f'На рентгенограмме {lr} лучезапястного сустава '
|
|
f'признаков перелома не выявлено.')
|
|
|
|
# Один найденный перелом
|
|
if num_frac == 1:
|
|
front_msg, side_msg = _format_single_fracture(labels[0])
|
|
front_count = 1 if front_msg else 0
|
|
else:
|
|
(front_msg, side_msg,
|
|
front_count, side_count) = _format_multiple_fractures(labels)
|
|
|
|
report_lines = []
|
|
if front_count:
|
|
report_lines.append(
|
|
f'На рентгенограмме {lr} лучезапястного сустава '
|
|
f'{front_msg}')
|
|
|
|
return ' '.join(report_lines).replace(' ', ' ')
|
|
|
|
|
|
def _make_conclusion(labels: list[str]) -> str:
|
|
num_frac = len(labels)
|
|
if num_frac == 0:
|
|
return 'Признаков перелома не выявлено.'
|
|
finds = {}
|
|
for label in labels:
|
|
name, part = label.split('_')[:2]
|
|
bone = bone_translator[part]
|
|
part = bone_parts_translator[name]
|
|
if bone in finds:
|
|
finds[bone].add(part)
|
|
else:
|
|
finds[bone] = {part}
|
|
find_texts = []
|
|
for bone in finds:
|
|
if 'метафиза' in finds[bone] and 'эпифиза' in finds[bone]:
|
|
finds[bone].remove('метафиза')
|
|
finds[bone].remove('эпифиза')
|
|
finds[bone].add('метаэпифиза')
|
|
|
|
if bone != 'лучевой или локтевой кости' or (
|
|
'лучевой кости' not in finds and
|
|
'локтевой кости' not in finds):
|
|
find_text = f'перелом '
|
|
find_text += ', '.join(sorted(finds[bone])) + f' {bone}'
|
|
find_texts.append(find_text)
|
|
conclusion = '; '.join(find_texts)
|
|
return conclusion.replace(' ', ' ').capitalize()
|
|
|
|
|
|
def _process_fractures(direct_img: Tensor,
|
|
combined: ProjectionSegments) -> tuple:
|
|
"""Обрабатываем боксы переломов, сортируем их и распределяем по
|
|
сегментам кости.
|
|
"""
|
|
frac_boxes, frac_scores = _get_fractions(direct_img)
|
|
frac_labels = [f'Находка {i + 1}' for i in range(len(frac_boxes))]
|
|
frac_boxes = frac_boxes[frac_boxes[:, 0].argsort()]
|
|
|
|
assigned, frac_boxes = _assign_fracs_to_bones(frac_boxes,
|
|
combined.bones_boxes,
|
|
combined.bones_labels)
|
|
hurt_bones = []
|
|
for bone_type, (cur_frac_boxes, cur_frac_labels) in assigned.items():
|
|
if bone_type == '4':
|
|
for label in cur_frac_labels:
|
|
hurt_bones.append(f'hand_4_{label}')
|
|
else:
|
|
mask = [lbl.endswith(f'_{bone_type}')
|
|
for lbl in combined.parts_labels]
|
|
mask_labels = np.array(combined.parts_labels)[mask]
|
|
hurt_bones += _assign_fracs_to_parts(cur_frac_boxes,
|
|
cur_frac_labels,
|
|
combined.parts_boxes[mask],
|
|
mask_labels)
|
|
return Fractures(frac_boxes, frac_scores, frac_labels, hurt_bones)
|
|
|
|
|
|
def _get_move_confidence(image: Tensor) -> float:
|
|
image = transform_move(image).unsqueeze(0).to(device)
|
|
|
|
with torch.no_grad():
|
|
outputs = model_move(image)[0]
|
|
if torch.argmax(outputs).item() == 1:
|
|
return 0
|
|
move_prob = torch.softmax(outputs, 0)[0]
|
|
return float(move_prob.item())
|
|
|
|
|
|
def _get_diastas_len(image: Tensor, parts: ProjectionSegments) -> int:
|
|
if '2' not in parts.bones_labels:
|
|
return 0
|
|
luch_box = parts.bones_boxes[parts.bones_labels.index('2')]
|
|
xmin, _, xmax = luch_box[:3]
|
|
bone_w = xmax - xmin
|
|
|
|
move_prob = _get_move_confidence(image)
|
|
|
|
coef = 0.025
|
|
move_len = (bone_w * move_prob * coef).item()
|
|
move_len = min(40, move_len)
|
|
|
|
return round(move_len)
|
|
|
|
|
|
def _visualize_detections(direct_img: Tensor, study_iuid: str,
|
|
laterality: Optional[str]) -> structs.Prediction:
|
|
"""Основная функция визуализации детекции с объединением боксов"""
|
|
if laterality in ['R', 'П']:
|
|
is_r = True
|
|
elif laterality in ['L', 'Л']:
|
|
is_r = False
|
|
else:
|
|
is_r = _check_is_right(direct_img)
|
|
|
|
bone_segments = _get_projection_bones(direct_img)
|
|
fracs = _process_fractures(direct_img, bone_segments)
|
|
|
|
report = _make_reports(fracs.bones, is_r)
|
|
conclusion = _make_conclusion(fracs.bones)
|
|
diastasis_mm = _get_diastas_len(direct_img, bone_segments)
|
|
diastasis_mm = diastasis_mm if fracs.labels else 0
|
|
|
|
img = cv2.cvtColor(direct_img[0].numpy(), cv2.COLOR_GRAY2RGB)
|
|
|
|
color = (255, 0, 0)
|
|
font = cv2.FONT_HERSHEY_COMPLEX
|
|
for box, label in zip(fracs.boxes.cpu().numpy().astype(int), fracs.labels):
|
|
cv2.rectangle(img, box[:2], box[2:], color, 2)
|
|
ytext = max(box[1] - 5, 0)
|
|
cv2.putText(img, label, (box[0], ytext), font, 0.5, color, 1)
|
|
|
|
cv2.imwrite(f"client/static/{study_iuid}.png", img)
|
|
|
|
is_fractured = bool(len(fracs.boxes))
|
|
if not is_fractured:
|
|
overall_probability = 0
|
|
else:
|
|
overall_probability = round(float(max(fracs.scores)*100))
|
|
|
|
properties = {
|
|
"Макс. величина диастаза отломков, мм": diastasis_mm
|
|
}
|
|
|
|
return structs.Prediction(overall_probability, is_fractured,
|
|
report, conclusion, img, properties)
|
|
|
|
|
|
def predict(input: structs.PredictorInput) -> structs.Prediction:
|
|
return _visualize_detections(input.image, input.study_uid,
|
|
input.laterality)
|