LORKT_AI_Service/service/reports.py

import os
import importlib
import base64
from zipfile import ZipFile
from datetime import datetime, timezone, timedelta
from pydicom.dataset import Dataset, validate_file_meta
from pydicom.uid import UID, generate_uid
import numpy as np
from service import sr_tags, preprocessor, structs

MODEL_VERSION = "1.0.0"


def _gen_seriesIUID(orig_seriesIUID, model, sr=False):
    if len(orig_seriesIUID) > 56:
        orig_seriesIUID = orig_seriesIUID[:56]
    match model:
        case "sinus":
            model_id = 1208
        case "wrist":
            model_id = 1249
        case "shoulder":
            model_id = 1250
    return f"{orig_seriesIUID}.{model_id}.{'2' if sr else '1'}"


def _to_sr_datetime(datetime):
    return datetime.strftime("%d.%m.%Y %H:%M")


def to_iso8601(datetime):
    msk_tz = timezone(timedelta(hours=3), name="MSK")
    date = datetime.astimezone(msk_tz).isoformat("T", "milliseconds")
    rindex_colon = date.rindex(':')
    return date[:rindex_colon] + date[rindex_colon+1:]


def _create_ds_base(meta_tags: structs.MetaTags, sop_class_uid: UID):
    meta_info = Dataset()
    sop_instance_uid = generate_uid()
    meta_info.MediaStorageSOPClassUID = sop_class_uid
    meta_info.MediaStorageSOPInstanceUID = sop_instance_uid
    meta_info.TransferSyntaxUID = UID('1.2.840.10008.1.2')

    ds = Dataset()
    ds.file_meta = meta_info
    validate_file_meta(ds.file_meta, enforce_standard=True)

    ds.SOPClassUID = sop_class_uid
    ds.InstitutionName = "LORKT"
    ds.StudyInstanceUID = meta_tags.study_iuid
    ds.SOPInstanceUID = sop_instance_uid
    if meta_tags.patient_id:
        ds.PatientID = meta_tags.patient_id
    if meta_tags.accession_number:
        ds.AccessionNumber = meta_tags.accession_number
    if meta_tags.issuer_of_patient_id:
        ds.IssuerOfPatientID = meta_tags.issuer_of_patient_id
    if meta_tags.filler_number:
        ds.FillerOrderNumberImagingServiceRequest = meta_tags.filler_number
    return ds


def _create_sr(meta_tags: structs.MetaTags, report: str, conclusion: str,
               process_end: datetime, model: str, username: str):
    sop_class_uid = UID('1.2.840.10008.5.1.4.1.1.88.33')
    ds = _create_ds_base(meta_tags, sop_class_uid)
    ds.SeriesInstanceUID = _gen_seriesIUID(meta_tags.series_iuid, model,
                                           sr=True)
    ds.Modality = "SR"
    ds.InstanceNumber = 1

    process_end = _to_sr_datetime(process_end)
    tags = sr_tags.tags_for_models[model]
    tags_and_texts = (
        ("Модальность", "РГ"),
        ("Область исследования", tags["Область исследования"]),
        ("Идентификатор исследования", meta_tags.study_iuid),
        ("Дата и время формирования заключения ИИ-сервисом", process_end),
        ("Предупреждение", "Заключение подготовлено программным обеспечением "\
                           "с применением технологий искусственного "\
                           "интеллекта"),
        ("Предупреждение", "В исследовательских целях"),
        ("Наименование сервиса", "ЛОР КТ"),
        ("Версия сервиса", MODEL_VERSION),
        ("Назначение сервиса", tags["Назначение сервиса"]),
        ("Технические данные", tags["Технические данные"]),
        ("Описание", report),
        ("Заключение", conclusion),
        ("Руководство пользователя", tags["Руководство пользователя"])
    )

    ds.SpecificCharacterSet = "ISO_IR 192"
    ds.add_new((0x0040, 0xa730), 'SQ', 
               [Dataset() for _ in range(len(tags_and_texts))])
    seq = ds.ContentSequence

    for i, (tag, text) in enumerate(tags_and_texts):
        seq[i].RelationshipType = "CONTAINS"
        seq[i].ValueType = "TEXT"
        seq[i].TextValue = text
        seq[i].add_new((0x0040, 0xa043), 'SQ', [Dataset()])
        name_seq = seq[i].ConceptNameCodeSequence[0]
        name_seq.CodeValue = "209001"
        name_seq.CodingSchemeDesignator = "99PMP"
        name_seq.CodeMeaning = tag

    save_path = f"data/{username}/sr/{meta_tags.study_iuid}.dcm"
    ds.save_as(save_path, implicit_vr=True, little_endian=True)
    return save_path


def _create_a_series(meta_tags: structs.MetaTags, img: np.ndarray, 
                     model: str, username: str):
    acquisition_date = datetime.now().strftime("%Y%m%d")
    acquisition_time = datetime.now().strftime("%H%M%S")
    series_uid = _gen_seriesIUID(meta_tags.series_iuid, model)
    sop_class_uid = UID('1.2.840.10008.5.1.4.1.1.7')
    ds = _create_ds_base(meta_tags, sop_class_uid)
    ds.SeriesInstanceUID = series_uid
    ds.InstanceNumber = 1
    ds.Modality = "DX"
    ds.SeriesDescription = "LORKT"
    ds.InstitutionName = "LORKT"
    ds.InstitutionalDepartmentName = MODEL_VERSION
    ds.AcquisitionDate = acquisition_date
    ds.AcquisitionTime = acquisition_time
    ds.OperatorsName = "AI"
    ds.PixelData = bytes(img)
    ds.Rows, ds.Columns = img.shape[:2]
    ds.BitsAllocated = 8
    ds.BitsStored = 8
    ds.HighBit = 7
    ds.SamplesPerPixel = 3
    ds.PhotometricInterpretation = "RGB"
    ds.PixelRepresentation = 0
    ds.PlanarConfiguration = 0
    save_path = f"data/{username}/additional_series/{meta_tags.study_iuid}.dcm"
    ds.save_as(save_path, implicit_vr=True, little_endian=True)
    return save_path


def _zip_reports(paths: list[str], username: str):
    study_uid = os.path.split(paths[0])[1]
    with ZipFile(f"data/{username}/reports/{study_uid.replace('.dcm', '.zip')}", 'w') as z:
        for path in paths:
            dir, id = os.path.split(path)
            dir = os.path.split(dir)[1]
            arcname = os.path.join(dir, id)
            z.write(path, arcname)


def make_reports(pathology: str, study_path: str, 
                 username: str) -> structs.Prediction:
    meta_tags, pred_input = preprocessor.prep_imgs(pathology, study_path)
    module = importlib.import_module("service.predictors." + pathology)
    predict_func = getattr(module, "predict")
    prediction = predict_func(pred_input)
    process_end = datetime.now()
    sr_path = _create_sr(meta_tags, prediction.report, prediction.conclusion, 
                         process_end, pathology, username)
    a_series_path = _create_a_series(meta_tags, prediction.image, pathology,
                                     username)
    _zip_reports([sr_path, a_series_path], username)
    prediction = prediction._replace(image=f"{meta_tags.study_iuid}.png")
    return prediction