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)