Базовый коммит

2025-12-03 16:50:06 +05:00
commit 939932eeb0
38 changed files with 2942 additions and 0 deletions
@@ -0,0 +1,222 @@
+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)