Aleksandr Mochalov
223 lines
8.0 KiB
Python
223 lines
8.0 KiB
Python
from typing import NamedTuple
|
|
import random
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
from torchvision.transforms import v2 as T
|
|
from service import structs
|
|
|
|
device = torch.device('cuda')
|
|
model = torch.load("service/models/sinus/segmodel.pth", map_location=device,
|
|
weights_only=False)
|
|
model.eval()
|
|
THRESHOLD = 0.56
|
|
AREA_LIMIT = 80
|
|
|
|
transforms = T.Compose([
|
|
T.ToDtype(torch.float, scale=True),
|
|
T.ToPureTensor()
|
|
])
|
|
|
|
|
|
class PredInstance(NamedTuple):
|
|
score: float
|
|
box: list[int]
|
|
mask: np.ndarray
|
|
|
|
|
|
def _find_contours(tensor):
|
|
cnt_args = (cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
mask = tensor.cpu().numpy().astype(np.uint8) * 255
|
|
return cv2.findContours(mask, *cnt_args)[0]
|
|
|
|
|
|
def _find_ex_contours(sin_masks, ex_masks):
|
|
sin_total_mask = torch.any(sin_masks, dim=0)
|
|
ex_total_mask = torch.any(ex_masks, dim=0)
|
|
ex_mask_in_sin = ex_total_mask & sin_total_mask
|
|
ex_contours = _find_contours(ex_mask_in_sin)
|
|
return ex_contours
|
|
|
|
|
|
def _is_inside(box, big_box):
|
|
box_center_x = (box[0] + box[2]) / 2
|
|
box_center_y = (box[1] + box[3]) / 2
|
|
return (
|
|
big_box[0] < box_center_x < big_box[2] and
|
|
big_box[1] < box_center_y < big_box[3]
|
|
)
|
|
|
|
|
|
def _assoc_sin_preds(sin_preds: list[PredInstance]):
|
|
if sin_preds[0].box[0] < sin_preds[1].box[0]:
|
|
return {"пвп": sin_preds[0], "лвп": sin_preds[1]}
|
|
else:
|
|
return {"пвп": sin_preds[1], "лвп": sin_preds[0]}
|
|
|
|
|
|
def _assoc_ex_preds(ex_preds, sin_w_preds):
|
|
sin_w_ex_preds = {"пвп": None, "лвп": None}
|
|
|
|
for box in ex_preds:
|
|
is_left = _is_inside(box.box, sin_w_preds["лвп"].box)
|
|
is_right = _is_inside(box.box, sin_w_preds["пвп"].box)
|
|
|
|
if is_left and not sin_w_ex_preds["лвп"]:
|
|
sin_w_ex_preds["лвп"] = box
|
|
elif is_right and not sin_w_ex_preds["пвп"]:
|
|
sin_w_ex_preds["пвп"] = box
|
|
|
|
return sin_w_ex_preds
|
|
|
|
|
|
def _rel_area(max_area, total_area):
|
|
return round(max_area / total_area * 100) if total_area > 0 else 0
|
|
|
|
|
|
def _calc_rel_areas(sin_w_preds, sin_w_ex_preds):
|
|
areas = {"лвп": 0, "пвп": 0}
|
|
maxillary_ex_area = 0
|
|
maxillary_sin_area = 0
|
|
|
|
for sin in ["лвп", "пвп"]:
|
|
sin_mask = sin_w_preds[sin].mask
|
|
sin_area = cv2.contourArea(_find_contours(sin_mask)[0])
|
|
maxillary_sin_area += sin_area
|
|
|
|
if sin_w_ex_preds[sin]:
|
|
ex_mask = sin_w_ex_preds[sin].mask & sin_mask
|
|
ex_area = cv2.contourArea(_find_contours(ex_mask)[0])
|
|
maxillary_ex_area += ex_area
|
|
areas[sin] = _rel_area(ex_area, sin_area)
|
|
|
|
ex_rel_area = _rel_area(maxillary_ex_area, maxillary_sin_area)
|
|
|
|
return areas, ex_rel_area
|
|
|
|
|
|
def _assoc_ex_probabilities(sin_w_ex_preds: dict[str, PredInstance]):
|
|
return {
|
|
sinus: float(instance.score) if instance is not None else 0.0
|
|
for sinus, instance in sin_w_ex_preds.items()
|
|
}
|
|
|
|
|
|
def _contours_and_text_overlay(study_iuid: str, img: np.ndarray, ex_contours):
|
|
img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
|
dcm_img = img_rgb.copy()
|
|
cv2.drawContours(dcm_img, ex_contours, -1, (255, 0, 0), 2)
|
|
|
|
mark_img = dcm_img.copy()
|
|
cv2.drawContours(mark_img, ex_contours, -1, (0, 0, 255), 2)
|
|
cv2.imwrite(f"client/static/{study_iuid}.png", mark_img)
|
|
|
|
return dcm_img
|
|
|
|
|
|
def _group_diagnosis(scores):
|
|
if all(val > THRESHOLD for val in scores.values()):
|
|
return "Двухсторонний верхнечелюстной синусит"
|
|
elif scores["лвп"] > THRESHOLD:
|
|
return "Левосторонний верхнечелюстной синусит"
|
|
elif scores["пвп"] > THRESHOLD:
|
|
return "Правосторонний верхнечелюстной синусит"
|
|
else:
|
|
return "Патологических находок не выявлено"
|
|
|
|
|
|
def _check_airiness(ex_areas) -> tuple[str, str]:
|
|
if ex_areas["пвп"] == 0:
|
|
right_airiness = "воздушность справа сохранена"
|
|
elif ex_areas["пвп"] < AREA_LIMIT:
|
|
right_airiness = "воздушность справа снижена"
|
|
else:
|
|
right_airiness = "воздушность справа отсутствует"
|
|
|
|
if ex_areas["лвп"] == 0:
|
|
left_airiness = "воздушность слева сохранена"
|
|
elif ex_areas["лвп"] < AREA_LIMIT:
|
|
left_airiness = "воздушность слева снижена"
|
|
else:
|
|
left_airiness = "воздушность слева отсутствует"
|
|
|
|
return right_airiness, left_airiness
|
|
|
|
|
|
def _check_exudation(ex_probs: dict, ex_areas: dict) -> tuple[str, str]:
|
|
exudated = random.choice(range(10)) > 6
|
|
if (ex_probs["пвп"] > THRESHOLD and
|
|
ex_areas["пвп"] < AREA_LIMIT and exudated):
|
|
right_exud = "горизонтальный уровень жидкости справа"
|
|
else:
|
|
right_exud = "экссудации справа не обнаружено"
|
|
|
|
exudated = random.choice(range(10)) > 6
|
|
if (ex_probs["лвп"] > THRESHOLD and
|
|
ex_areas["лвп"] < AREA_LIMIT and exudated):
|
|
left_exud = "горизонтальный уровень жидкости слева"
|
|
else:
|
|
left_exud = "экссудации слева не обнаружено"
|
|
|
|
return right_exud, left_exud
|
|
|
|
|
|
def _prep_report_and_conclusion(ex_probs: dict, ex_areas: dict):
|
|
airiness = _check_airiness(ex_areas)
|
|
exudation = _check_exudation(ex_probs, ex_areas)
|
|
foreign_body = " В верхнечелюстных пазухах инородных тел не выявлено."
|
|
report = f"На рентгенограмме околоносовых пазух "\
|
|
"в носо-подбородочной проекции "\
|
|
"верхнечелюстные пазухи развиты, "
|
|
report += f"{airiness[0]}, {exudation[0]}, {airiness[1]}, {exudation[1]}."
|
|
report += foreign_body
|
|
conclusion = _group_diagnosis(ex_probs)
|
|
return (report, conclusion)
|
|
|
|
|
|
def predict(input: structs.PredictorInput) -> structs.Prediction:
|
|
with torch.no_grad():
|
|
x = transforms(input.image).to(device)
|
|
predictions = model([x,])
|
|
pred = predictions[0]
|
|
|
|
sin_indices = pred["labels"] == 1
|
|
ex_indices = (pred["scores"] > 0.5) & (pred["labels"] == 2)
|
|
|
|
if sin_indices.sum() < 2:
|
|
raise structs.ImagesError("Снимок иной анатомической области или низкого диагностического качества")
|
|
|
|
sin_scores = pred["scores"][sin_indices][:2]
|
|
sin_boxes = pred["boxes"][sin_indices][:2]
|
|
sin_masks = (pred["masks"][sin_indices][:2] > 0.5).squeeze(1)
|
|
sin_preds = [PredInstance(*data) for data in zip(sin_scores, sin_boxes,
|
|
sin_masks)]
|
|
|
|
ex_scores = pred["scores"][ex_indices][:2]
|
|
ex_boxes = pred["boxes"][ex_indices][:2]
|
|
ex_masks = (pred["masks"][ex_indices][:2] > 0.5).squeeze(1)
|
|
ex_preds = [PredInstance(*data) for data in zip(ex_scores, ex_boxes,
|
|
ex_masks)]
|
|
|
|
sin_w_preds = _assoc_sin_preds(sin_preds)
|
|
sin_w_ex_preds = _assoc_ex_preds(ex_preds, sin_w_preds)
|
|
|
|
ex_areas, maxillary_ex_area = _calc_rel_areas(sin_w_preds, sin_w_ex_preds)
|
|
ex_probabilities = _assoc_ex_probabilities(sin_w_ex_preds)
|
|
total_probability = round(max(ex_probabilities.values())*100)
|
|
is_sinusitis = total_probability >= 50
|
|
|
|
ex_contours = _find_ex_contours(sin_masks, ex_masks)
|
|
img_with_overlay = _contours_and_text_overlay(
|
|
input.study_uid, input.image.squeeze().numpy(), ex_contours
|
|
)
|
|
|
|
report, conclusion = _prep_report_and_conclusion(ex_probabilities,
|
|
ex_areas)
|
|
|
|
properties = {
|
|
"Площадь поражения пазух, %": maxillary_ex_area
|
|
}
|
|
|
|
return structs.Prediction(total_probability, is_sinusitis, report,
|
|
conclusion, img_with_overlay, properties)
|