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)