| import albumentations as A |
| import torch |
| import torch.nn as nn |
|
|
| from numpy.typing import NDArray |
| from transformers import PreTrainedModel |
| from timm import create_model |
| from typing import Optional |
| from .configuration import CXRConfig |
| from .unet import UnetDecoder, SegmentationHead |
|
|
| _PYDICOM_AVAILABLE = False |
| try: |
| from pydicom import dcmread |
| from pydicom.pixels import apply_voi_lut |
|
|
| _PYDICOM_AVAILABLE = True |
| except ModuleNotFoundError: |
| pass |
|
|
|
|
| class CXRModel(PreTrainedModel): |
| config_class = CXRConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.encoder = create_model( |
| model_name=config.backbone, |
| features_only=True, |
| pretrained=False, |
| in_chans=config.in_chans, |
| ) |
| self.decoder = UnetDecoder( |
| decoder_n_blocks=config.decoder_n_blocks, |
| decoder_channels=config.decoder_channels, |
| encoder_channels=config.encoder_channels, |
| decoder_center_block=config.decoder_center_block, |
| decoder_norm_layer=config.decoder_norm_layer, |
| decoder_attention_type=config.decoder_attention_type, |
| ) |
| self.img_size = config.img_size |
| self.segmentation_head = SegmentationHead( |
| in_channels=config.decoder_channels[-1], |
| out_channels=config.seg_num_classes, |
| size=self.img_size, |
| ) |
| self.pooling = nn.AdaptiveAvgPool2d(1) |
| self.dropout = nn.Dropout(p=config.cls_dropout) |
| self.classifier = nn.Linear(config.feature_dim, config.cls_num_classes) |
|
|
| def normalize(self, x: torch.Tensor) -> torch.Tensor: |
| |
| mini, maxi = 0.0, 255.0 |
| x = (x - mini) / (maxi - mini) |
| x = (x - 0.5) * 2.0 |
| return x |
|
|
| @staticmethod |
| def load_image_from_dicom(path: str) -> Optional[NDArray]: |
| if not _PYDICOM_AVAILABLE: |
| print("`pydicom` is not installed, returning None ...") |
| return None |
| dicom = dcmread(path) |
| arr = apply_voi_lut(dicom.pixel_array, dicom) |
| if dicom.PhotometricInterpretation == "MONOCHROME1": |
| |
| arr = arr.max() - arr |
|
|
| arr = arr - arr.min() |
| arr = arr / arr.max() |
| arr = (arr * 255).astype("uint8") |
| return arr |
|
|
| def preprocess(self, x: NDArray) -> NDArray: |
| x = A.Resize(self.img_size[0], self.img_size[1], p=1)(image=x)["image"] |
| return x |
|
|
| def forward(self, x: torch.Tensor, return_logits: bool = False) -> torch.Tensor: |
| x = self.normalize(x) |
| features = self.encoder(x) |
| decoder_output = self.decoder(features) |
| logits = self.segmentation_head(decoder_output[-1]) |
| b, n = features[-1].shape[:2] |
| features = self.pooling(features[-1]).reshape(b, n) |
| features = self.dropout(features) |
| cls_logits = self.classifier(features) |
| out = { |
| "mask": logits, |
| "age": cls_logits[:, 0].unsqueeze(1), |
| "view": cls_logits[:, 1:4], |
| "female": cls_logits[:, 4].unsqueeze(1), |
| } |
| if return_logits: |
| return out |
| out["mask"] = out["mask"].softmax(1) |
| out["view"] = out["view"].softmax(1) |
| out["female"] = out["female"].sigmoid() |
| return out |
|
|
