Aleksandr Mochalov
86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
from typing import NamedTuple
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
import pydicom
|
|
from service.structs import PredictorInput, TagError, MetaTags
|
|
|
|
|
|
class InputImage(NamedTuple):
|
|
img: np.ndarray
|
|
ww: int
|
|
wc: int
|
|
uid: str
|
|
color_inversion: bool
|
|
|
|
|
|
def _get_ww_wc(study_ww, study_wc):
|
|
if type(study_ww) is pydicom.valuerep.DSfloat:
|
|
return int(study_ww), int(study_wc)
|
|
return int(study_ww[0]), int(study_wc[0])
|
|
|
|
|
|
def _is_color_inverted(study):
|
|
return study.PhotometricInterpretation == "MONOCHROME1"
|
|
|
|
|
|
def _prep_img(img_pack: InputImage) -> torch.Tensor:
|
|
img = img_pack.img.astype(np.float32)
|
|
lower_bound = img_pack.wc - 0.5 * img_pack.ww
|
|
upper_bound = img_pack.wc + 0.5 * img_pack.ww
|
|
img = np.clip(img, lower_bound, upper_bound)
|
|
img = ((img - lower_bound) / img_pack.ww) * 255
|
|
img = img.astype(np.uint8)
|
|
if img_pack.color_inversion:
|
|
img = 255 - img
|
|
return torch.from_numpy(img[None, ...])
|
|
|
|
|
|
def _check_required_tags(study: pydicom.FileDataset, pathology: str):
|
|
invalid_tags = []
|
|
if not (hasattr(study, "PhotometricInterpretation") and
|
|
study.PhotometricInterpretation in ("MONOCHROME1", "MONOCHROME2")):
|
|
invalid_tags.append("PhotometricInterpretation")
|
|
if pathology in ("shoulder", "wrist"):
|
|
if not (hasattr(study, "PixelSpacing") or
|
|
hasattr(study, "ImagerPixelSpacing")):
|
|
invalid_tags.append("PixelSpacing/ImagerPixelSpacing")
|
|
for tag in ("WindowWidth", "WindowCenter"):
|
|
if not hasattr(study, tag) or getattr(study, tag) == "":
|
|
invalid_tags.append(tag)
|
|
if len(invalid_tags) == 1:
|
|
raise TagError(f"DICOM тег {invalid_tags[0]} не заполнен, "\
|
|
"либо имеет некорректное значение")
|
|
elif len(invalid_tags) > 1:
|
|
raise TagError(f"DICOM теги {", ".join(invalid_tags)} не заполнены, "\
|
|
"либо имеют некорректные значения")
|
|
|
|
|
|
def _get_meta_tags(study: pydicom.FileDataset):
|
|
return MetaTags(study.StudyInstanceUID, study.SeriesInstanceUID,
|
|
getattr(study, "PatientID", None),
|
|
getattr(study, "AccessionNumber", None),
|
|
getattr(study, "IssuerOfPatientID", None),
|
|
getattr(study, "FillerOrderNumberImagingServiceRequest",
|
|
None))
|
|
|
|
|
|
def _get_px_size(study: pydicom.FileDataset):
|
|
if hasattr(study, "PixelSpacing"):
|
|
return study.PixelSpacing
|
|
if hasattr(study, "ImagerPixelSpacing"):
|
|
return study.ImagerPixelSpacing
|
|
return None, None
|
|
|
|
|
|
def prep_imgs(pathology: str, study_path: str) -> tuple[MetaTags, PredictorInput]:
|
|
instance = pydicom.dcmread(study_path, force=True)
|
|
_check_required_tags(instance, pathology)
|
|
meta_tags = _get_meta_tags(instance)
|
|
ww, wc = _get_ww_wc(instance.WindowWidth, instance.WindowCenter)
|
|
clr_inverted = _is_color_inverted(instance)
|
|
img = InputImage(instance.pixel_array, ww, wc, meta_tags.study_iuid, clr_inverted)
|
|
pred_input = PredictorInput(meta_tags.study_iuid, _prep_img(img),
|
|
getattr(instance, "ImageLaterality", None),
|
|
*_get_px_size(instance))
|
|
return meta_tags, pred_input |