import warnings from typing import Optional, NamedTuple import cv2 import numpy as np import torch from torch import Tensor import torchvision.ops as ops from torchvision.transforms import v2 as T from skimage.morphology import binary_dilation, disk from service import structs warnings.filterwarnings('ignore', category=UserWarning) device = torch.device('cuda') models_root = "service/models/shoulder" model_frac = torch.load(f'{models_root}/frac_model.pth', weights_only=False).to(device) model_lr = torch.load(f'{models_root}/lr_model.pth', weights_only=False).to(device) model_parts = torch.load(f'{models_root}/parts_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_lr.eval() model_parts.eval() model_move.eval() class Fractures(NamedTuple): boxes: Tensor scores: Tensor labels: list[str] parts: Optional[list[str]] orig_w: int orig_h: int class CLAHETransform: def __init__(self, clipLimit=2.0, tileGridSize=(8, 8)): self.clipLimit = clipLimit self.tileGridSize = tileGridSize self.clahe = cv2.createCLAHE(clipLimit=self.clipLimit, tileGridSize=self.tileGridSize) def __call__(self, img, target=None): 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) cl_img_tensor = torch.from_numpy(cl_img).unsqueeze(0) cl_img_tensor = cl_img_tensor.to(img.device) return cl_img_tensor transform_parts = T.Compose([ T.Resize((256, 256)), T.Grayscale(num_output_channels=1), CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)), T.ToDtype(torch.float, scale=True), T.ToTensor(), T.Normalize(mean=[0.0773] * 3, std=[0.0516] * 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.40526121854782104], std=[0.23242981731891632]) ]) transform_lr = T.Compose([ T.Resize((256, 256)), CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)), T.ToDtype(torch.float, scale=True), T.ToPureTensor(), T.Normalize(mean=[0.0773] * 3, std=[0.0516] * 3) ]) transform_move = transform_test = T.Compose([ T.Resize((224, 224)), T.Grayscale(num_output_channels=1), CLAHETransform(clipLimit=2.0, tileGridSize=(8, 8)), T.ToDtype(torch.float, scale=True), T.ToTensor(), T.Normalize(mean=[0.4172] * 3, std=[0.2612] * 3) ]) parts_map = { 0: 'акромиального отростка', 1: 'клювовидного отростка', 2: 'плечевой кости', 3: 'суставной впадины', 4: 'тела', 5: 'шейки', 6: 'Фон', 7: 'головки', 8: 'анатомической шейки', 9: 'хирургической шейки', 10: 'диафиза' } parts2bones = { 0: 'Лопатка', 1: 'Лопатка', 2: 'Плечевая кость', 3: 'Лопатка', 4: 'Лопатка', 5: 'Лопатка', 6: 'Фон', 7: 'Плечевая кость', 8: 'Плечевая кость', 9: 'Плечевая кость', 10: 'Плечевая кость' } def _convert_bboxes(bboxes: Tensor, orig_width: int, orig_height: int, transformed_width=512, transformed_height=512) -> Tensor: """Масштабирует координаты боксов обратно к размерам исходного изображения.""" conv_bboxes = bboxes.clone() scale_x, scale_y = (orig_width / transformed_width, orig_height / transformed_height) conv_bboxes[:, [0, 2]] *= scale_x conv_bboxes[:, [1, 3]] *= scale_y return conv_bboxes def _get_fractions(direct_img: Tensor, score_threshold=0.07) -> Fractures: """Получение переломов""" original = direct_img.clone() 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.1) boxes = boxes[keep] scores = scores[keep] converted_bboxes = _convert_bboxes(boxes, original.shape[2], original.shape[1]) converted_bboxes = converted_bboxes[converted_bboxes[:, 0].argsort()] fractures = Fractures(converted_bboxes, scores, [], None, original.shape[2], original.shape[1]) return fractures 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 _smooth_segmentation_mask(mask: np.ndarray, kernel_size=5, min_area=10) -> np.ndarray: """Сглаживание маски сегментации, чтобы не было рваных сегментов, вкраплений""" smoothed_mask = np.zeros_like(mask) n_classes = int(mask.max()) + 1 for cls in range(n_classes): class_mask = (mask == cls).astype(np.uint8) closed = cv2.morphologyEx(class_mask, cv2.MORPH_CLOSE, np.ones((kernel_size, kernel_size), np.uint8)) opened = cv2.morphologyEx(closed, cv2.MORPH_OPEN, np.ones((kernel_size, kernel_size), np.uint8)) num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(opened, connectivity=4) clean = np.zeros_like(opened) for i in range(1, num_labels): if stats[i, cv2.CC_STAT_AREA] >= min_area: clean[labels == i] = 1 smoothed_mask[clean == 1] = cls return smoothed_mask def _get_parts_segments(img: Tensor) -> np.ndarray: img = img.repeat(3, 1, 1) image = transform_parts(img).unsqueeze(0).to(device) with torch.no_grad(): outputs = model_parts(image)[0] output = outputs.detach().cpu().numpy() predicted_classes = np.argmax(output, axis=0) if np.sum(predicted_classes == 6) > 64000: raise structs.ImagesError("Снимок иной анатомической области или низкого диагностического качества") predicted_classes = _smooth_segmentation_mask(predicted_classes) return predicted_classes def _compute_pca_direction(coords: np.ndarray) -> np.ndarray: mean = coords.mean(axis=0) centered = coords - mean cov = np.cov(centered, rowvar=False) eigvals, eigvecs = np.linalg.eigh(cov) principal = eigvecs[:, np.argmax(eigvals)] return principal / np.linalg.norm(principal) def _rotate_vector(vec: np.ndarray, angle_deg: float) -> np.ndarray: theta = np.deg2rad(angle_deg) rot = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) return rot.dot(vec) def _get_segment2_coords(mask: np.ndarray, class_idx=2) -> np.ndarray: ys, xs = np.where(mask == class_idx) return np.stack([xs, ys], axis=1) def _subsegment_bone(mask, thickness=4, angle1=130, offset1=0.02, angle2=90, offset2=0.25): """Разбиение плечевой кости на микро-сегменты""" mask = mask.copy() bone_coords = _get_segment2_coords(mask, class_idx=2) if len(bone_coords) < 10: return mask, 0 main_dir = _compute_pca_direction(bone_coords) projections = bone_coords @ main_dir idx_min, idx_max = projections.argmin(), projections.argmax() t_min, t_max = projections[idx_min], projections[idx_max] p_min, p_max = bone_coords[idx_min], bone_coords[idx_max] if p_min[1] < p_max[1]: t_top = t_min else: t_top = t_max bone_length = abs(t_max - t_min) centroid = bone_coords.mean(axis=0) centroid_proj = centroid.dot(main_dir) H, W = mask.shape y_grid, x_grid = np.mgrid[0:H, 0:W] pix = np.stack([x_grid.ravel(), y_grid.ravel()], axis=1) seg2_mask_flat = (mask.ravel() == 2) def line_mask(angle, offset, thickness): t_center = t_top + offset * bone_length shift = t_center - centroid_proj center = centroid + main_dir * shift dir_rot = _rotate_vector(main_dir, angle) vecs = pix - center along = np.dot(vecs, dir_rot) ortho_vecs = vecs - np.outer(along, dir_rot) dists = np.linalg.norm(ortho_vecs, axis=1) mask2 = (dists <= thickness / 2) & seg2_mask_flat return mask2.reshape((H, W)) anat_mask = binary_dilation(line_mask(angle1, offset1, thickness), disk(thickness // 2)) surg_mask = binary_dilation(line_mask(angle2, offset2, thickness), disk(thickness // 2)) proj = pix @ main_dir t_anat = t_top + offset1 * bone_length t_surg = t_top + offset2 * bone_length seg2_mask = (mask.ravel() == 2) diaf_mask = (proj > t_surg) & seg2_mask diaf_mask = diaf_mask.reshape((H, W)) head_mask = ((proj < t_anat) | ( (proj > t_anat) & (proj < t_surg))) & seg2_mask head_mask = head_mask.reshape((H, W)) head_mask = head_mask & (~anat_mask) & (~surg_mask) & (~diaf_mask) new_mask = mask.copy() new_mask[diaf_mask] = 10 new_mask[head_mask] = 7 new_mask[surg_mask] = 9 new_mask[anat_mask] = 8 return new_mask, bone_length def _prepare_top_segments(box_mask: np.ndarray, weights: np.ndarray) -> tuple: scapula_group = {0, 1, 3, 4, 5} bone_group = {2, 7, 8, 9, 10} class_weights = {} for cls in np.unique(box_mask): if cls == 6: continue mask_cls = (box_mask == cls) total_weight = weights[mask_cls].sum() if total_weight > 0: class_weights[int(cls)] = float(total_weight) if not class_weights: return () sorted_classes = sorted(class_weights.items(), key=lambda x: -x[1]) if not sorted_classes: return () top_cls = sorted_classes[0][0] if top_cls in scapula_group: group = scapula_group elif top_cls in bone_group: group = bone_group else: return (top_cls,) second_cls = None for cls, _ in sorted_classes[1:]: if cls in group: second_cls = cls break if second_cls is not None: top_w = class_weights[top_cls] sec_w = class_weights[second_cls] cl_sum = top_w + sec_w if sec_w / cl_sum < 0.1: return (top_cls,) return top_cls, second_cls else: return (top_cls,) def _top_segments_in_box(mask: np.ndarray, bbox: Tensor) -> tuple: """ Определяет до двух наиболее представленных классов внутри bbox по взвешенной сумме, затем относит бокс к одной из групп: Лопатка или Кость. """ bbox = bbox.cpu().numpy().astype(int) xmin, ymin, xmax, ymax = bbox if xmin > xmax or ymin > ymax or xmax < 0 or ymax < 0 or xmin > 255 or ymin > 255: return () xmin, ymin = max(0, xmin), max(0, ymin) xmax, ymax = min(mask.shape[1] - 1, xmax), min(mask.shape[0] - 1, ymax) box_mask = mask[ymin:ymax + 1, xmin:xmax + 1] h, w = box_mask.shape ys, xs = np.mgrid[0:h, 0:w] center_y, center_x = (h - 1) / 2, (w - 1) / 2 dists = np.sqrt((ys - center_y) ** 2 + (xs - center_x) ** 2) max_dist = dists.max() if dists.max() > 0 else 1.0 weights = np.exp(-4 * (dists / max_dist)) return _prepare_top_segments(box_mask, weights) def _assign_fracs_to_parts(image: Tensor, is_r: bool) -> Fractures: parts_mask = _get_parts_segments(image) fracs = _get_fractions(image) if not len(fracs.boxes): return fracs, 0 conv_boxes = _convert_bboxes(fracs.boxes, 256, 256, fracs.orig_w, fracs.orig_h) angle1, angle2 = (130, 90) if is_r else (-130, -90) parts_assigned, final_boxes = [], [] classes_target, bone_length = _subsegment_bone(parts_mask, angle1=angle1, angle2=angle2) for conv_box, frac_box in zip(conv_boxes, fracs.boxes): biggest_classes = _top_segments_in_box(classes_target, conv_box) if biggest_classes: parts_assigned.append(biggest_classes) final_boxes.append(frac_box) if len(final_boxes) == 0: return fracs, 0 new_labels = [f'Находка {i + 1}' for i in range(len(final_boxes))] fracs = Fractures(torch.stack(final_boxes), fracs.scores, new_labels, parts_assigned, fracs.orig_w, fracs.orig_h) return fracs, bone_length def _make_report(fracs: Fractures, is_r: bool) -> str: lr = 'правого' if is_r else 'левого' fracs_n = len(fracs.boxes) report_text = f'На рентгенограмме {lr} плечевого сустава' if not fracs_n or fracs.parts is None: report_text += ' переломов не выявлено.' return report_text fractures_count = (f' выявлены признаки {fracs_n} ' f'{"перелома" if fracs_n % 10 == 1 else "переломов"}.') report_text += fractures_count for label, parts in zip(fracs.labels, fracs.parts): parts_text = f' {label}: перелом ' parts_text += ', '.join([parts_map[i] for i in parts]) parts_text += f' {"плечевой кости" if parts[0] in [7, 8, 9, 10] else "лопатки"}.' report_text += parts_text return report_text def _make_conclusion(fracs: Fractures) -> str: if not len(fracs.boxes) or fracs.parts is None: return 'Признаков перелома не выявлено.' finds = {} for parts in fracs.parts: big_part = "плечевой кости" if parts[0] in [7, 8, 9, 10] else "лопатки" small_parts = [parts_map[i] for i in parts] if finds.get(big_part): finds[big_part].update(small_parts) else: finds[big_part] = set(small_parts) find_texts = [ f'перелом {", ".join(sorted(parts))} {bone}' for bone, parts in finds.items() ] conclusion = '; '.join(find_texts) + '.' return conclusion.replace(' ', ' ').capitalize() def _get_move_confidence(image): image = image.repeat(3, 1, 1) 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_move_len(image: Tensor, bone_length): move_prob = _get_move_confidence(image)-0.5 coef = 0.0650 move_len = bone_length * move_prob * coef move_len = max(0, 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) img = cv2.cvtColor(direct_img[0].numpy(), cv2.COLOR_GRAY2RGB) font = cv2.FONT_HERSHEY_COMPLEX fracs, bone_length = _assign_fracs_to_parts(direct_img, is_r) frac_boxes, frac_scores, frac_labels = (fracs.boxes, fracs.scores, fracs.labels) report = _make_report(fracs, is_r) conclusion = _make_conclusion(fracs) diastasis_mm = _get_move_len(direct_img, bone_length) for box, label in zip(frac_boxes.cpu().numpy(), frac_labels): xmin, ymin, xmax, ymax = box.astype(int) cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(255, 0, 0), thickness=2) cv2.putText(img, label, (xmin, max(ymin - 5, 0)), font, 0.5, (255, 0, 0), thickness=1, lineType=cv2.LINE_AA) cv2.imwrite(f"client/static/{study_iuid}.png", img) is_fractured = bool(len(frac_boxes)) if not is_fractured: overall_probability = 0 else: overall_probability = round(float(max(frac_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)