Базовый коммит
This commit is contained in:
Aleksandr Mochalov
@@ -0,0 +1,494 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user