Segmentation output writers¶

@dataclass
class SegWriter:
    """
    Helper class to manage DICOM Segmentation (SEG) and RT-Struct export.

    This class handles the conversion of prediction masks into standardized
    DICOM formats, including metadata management for algorithms and segments.

    Attributes:
        algorithm_name (str): Name of the segmentation algorithm.
        segment_names (list[str | dict[str, str]], optional): Names or dicts
            describing the segments.
        segment_descriptions (list[hd.seg.SegmentDescription], optional):
            High-level segment descriptions.
        algorithm_version (str): Version string of the algorithm.
        algorithm_family (Code): DICOM code for the algorithm family.
        algorithm_type (SegmentAlgorithmTypeValues): DICOM algorithm type.
        instance_number (int): DICOM instance number.
        series_number (int): DICOM series number.
        manufacturer (str): Manufacturer name for DICOM metadata.
        manufacturer_model_name (str): Model name for DICOM metadata.
        series_description (str): Description for the DICOM series.
        clinical_trial_series_id (str): Clinical trial series ID.
        clinical_trial_time_point_id (str): Clinical trial time point ID.
        body_part_examined (str): Body part examined for DICOM metadata.
        validate (bool): Whether to perform validation on init.
    """

    algorithm_name: str
    segment_names: list[str | dict[str, str]] | None = None
    segment_descriptions: list[hd.seg.SegmentDescription] | None = None
    algorithm_version: str = "v1.0"
    algorithm_family: pydicom.sr.coding.Code = (
        codes.CID7162.ArtificialIntelligence
    )
    algorithm_type: hd.seg.SegmentAlgorithmTypeValues = (
        hd.seg.SegmentAlgorithmTypeValues.AUTOMATIC
    )
    instance_number: int = 1
    series_number: int = 999
    manufacturer: str = "Algorithm"
    manufacturer_model_name: str = "AlgorithmModel"
    series_description: str = "Segmentation"
    clinical_trial_series_id: str = "1"
    clinical_trial_time_point_id: str = "1"
    body_part_examined: str = "BODY"
    validate: bool = False

    def __post_init__(self):
        if self.segment_descriptions is None and self.segment_names is None:
            raise ValueError(
                "Either segment_descriptions or segment_names must be provided"
            )
        category_concepts = CATEGORY_CONCEPTS[DEFAULT_SEGMENT_SCHEME]
        self.algorithm_identification = hd.AlgorithmIdentificationSequence(
            name=self.algorithm_name,
            version=self.algorithm_version,
            family=self.algorithm_family,
        )
        if self.segment_names is None:
            return
        if self.segment_descriptions is None:
            self.segment_descriptions = []
        for i, segment in enumerate(self.segment_names):
            type_code, segment_dict = get_segment_type_code(segment, i)
            category_code = [
                category_concepts[k]
                for k in category_concepts
                if category_concepts[k].value == segment_dict["category_number"]
            ][0]
            random_rgb_colour = random_color_generator()
            segment_description = hd.seg.SegmentDescription(
                segment_number=segment_dict["number"],
                segment_label=segment_dict["label"],
                segmented_property_category=category_code,
                segmented_property_type=type_code,
                algorithm_type=self.algorithm_type,
                algorithm_identification=self.algorithm_identification,
                tracking_uid=hd.UID(),
                tracking_id=segment_dict["tracking_id"],
            )
            segment_description.RecommendedDisplayCIELabValue = rgb_to_cielab(
                random_rgb_colour
            )
            if segment_dict["laterality"] is not None:
                lat_code = LATERALITY_CODING[segment_dict["scheme"]][
                    segment_dict["laterality"]
                ]
                segment_description.SegmentedPropertyTypeModifierCodeSequence = [
                    lat_code
                ]
            csi = CODING_SCHEME_INFORMATION[segment_dict["scheme"]]
            for (
                prop_type
            ) in segment_description.SegmentedPropertyTypeCodeSequence:
                for k in csi:
                    prop_type[k] = DataElement(
                        tag=k,
                        VR=CODING_SCHEME_INFORMATION_VR[k],
                        value=csi[k],
                    )
            for (
                prop_cat
            ) in segment_description.SegmentedPropertyCategoryCodeSequence:
                for k in csi:
                    prop_cat[k] = DataElement(
                        tag=k,
                        VR=CODING_SCHEME_INFORMATION_VR[k],
                        value=csi[k],
                    )
            self.segment_descriptions.append(segment_description)

        if self.validate is True:
            for i, segment_description in enumerate(self.segment_descriptions):
                logger.info(
                    "Segment %s: %s, %s, class_idx=%i",
                    segment_description.segment_number,
                    segment_description.segment_label,
                    segment_description.segmented_property_type.meaning,
                    i,
                )

    def to_array_if_necessary(
        self, mask: np.ndarray | sitk.Image
    ) -> np.ndarray:
        """
        Converts a mask to a numpy array if the mask is a sitk.Image.

        Args:
            mask (np.ndarray | sitk.Image): the mask to convert.

        Returns:
            np.ndarray: the mask as a numpy array.
        """
        if isinstance(mask, sitk.Image):
            mask = sitk.GetArrayFromImage(mask)
        return mask

    def make_compliant(self, f: pydicom.Dataset):
        """
        Tries to smooth some potential missing fields.

        Args:
            f (pydicom.Dataset): a pydicom Dataset.

        Returns:
            A compliant pydicom.Dataset.
        """
        if hasattr(f, "PatientSex") is False:
            f.PatientSex = "O"
        if hasattr(f, "AccessionNumber") is False:
            f.AccessionNumber = ""
        if hasattr(f, "StudyID") is False:
            f.StudyID = ""
        return f

    def write_dicom_seg(
        self,
        mask_array: np.ndarray | sitk.Image,
        source_files: list[str],
        output_path: str,
        is_fractional: bool = False,
        is_fractional_compliant: bool = False,
        class_idx: int | None = None,
    ):
        """
        Writes a DICOM segmentation file.

        Args:
            mask_array (np.ndarray | sitk.Image): the mask array or sitk image.
            source_files (list[str]): the list of DICOM source files (as returned
                by ``nnunet_serve.utils.read_dicom_as_sitk``).
            output_path (str): the output path.
            is_fractional (bool, optional): whether the mask is fractional.
                Defaults to False.
            is_fractional_compliant (bool, optional): whether the probability mask
                should be converted to a map with ``N_FRACTIONAL`` labels,
                each corresponding to a percentage.
            class_idx (int | None, optional): index used for selecting specific
                segment descriptions when saving probability maps. Only used when
                either ``is_fractional==True`` or
                ``is_fractional_compliant==True``.
        """
        mask_array = self.to_array_if_necessary(mask_array)
        sorted_source_files = sort_dicom_slices(list(source_files))
        any_fractional = is_fractional or is_fractional_compliant
        if sorted_source_files != list(source_files):
            idx_map = {p: i for i, p in enumerate(source_files)}
            try:
                order = [idx_map[p] for p in sorted_source_files]
            except KeyError:
                order = None
            if order is not None:
                mask_array = mask_array[order, ...]
            source_files = sorted_source_files
        # adjust array size and segment descriptions to the strictly necessary
        labels = np.unique(mask_array)
        labels = labels[labels > 0]
        if len(labels) == 0:
            logger.warning("Mask is empty")
            return "empty"
        segment_descriptions = []
        if any_fractional is False:
            label_dict = {label: i + 1 for i, label in enumerate(labels)}
            label_dict[0] = 0
            mask_array = np.vectorize(label_dict.get)(mask_array)
            for i, label in enumerate(labels.astype(int)):
                seg_d = deepcopy(self.segment_descriptions[label - 1])
                seg_d.SegmentNumber = i + 1
                segment_descriptions.append(seg_d)
            if len(mask_array.shape) != 4:
                mask_array = one_hot_encode(
                    mask_array, len(segment_descriptions)
                )
        else:
            if class_idx is not None:
                if isinstance(class_idx, list):
                    if len(class_idx) == 1:
                        class_idx = class_idx[0]
                    else:
                        raise ValueError(
                            f"class_idx should be int or a list with length 1 (is {class_idx})"
                        )
                segment_descriptions = [
                    self.segment_descriptions[class_idx - 1]
                ]
            else:
                segment_descriptions = self.segment_descriptions
            if mask_array.ndim == 3:
                mask_array = mask_array[..., None]
        image_datasets = [
            self.make_compliant(hd.imread(str(f))) for f in source_files
        ]

        first = image_datasets[0]
        rows, cols = int(getattr(first, "Rows")), int(getattr(first, "Columns"))
        frames = len(image_datasets)
        if (frames, rows, cols) != mask_array.shape[:-1]:
            raise Exception(
                f"Mask shape {mask_array.shape} does not match image shape {frames}x{rows}x{cols}"
            )
        if (
            mask_array.shape[-1] != len(segment_descriptions)
            and any_fractional is False
        ):
            # this check makes no sense for fractional outputs!
            raise Exception(
                f"Mask shape {mask_array.shape} does not match number of segments {len(segment_descriptions)}"
            )

        # Create the Segmentation instance
        label_meanings = []
        for s in segment_descriptions:
            ss = s.SegmentedPropertyTypeCodeSequence[0]
            if isinstance(ss, pydicom.Dataset):
                meaning = ss.CodeMeaning
            else:
                meaning = ss.meaning
            label_meanings.append(meaning.lower().replace(" structure of", ""))
        seg_series_description = "Seg of " + ", ".join(label_meanings)
        if len(seg_series_description) > 64:
            seg_series_description = seg_series_description[:61] + "..."
        if is_fractional_compliant:
            if len(segment_descriptions) > 1:
                logger.warning(
                    "Skipping saving the fractional compliant mask because len(segment_descriptions) > 1"
                )
                return "skipped"
            seg_type = hd.seg.SegmentationTypeValues.BINARY
            logger.info("Converting mask to pseudo-fractional DICOM seg")
            sd = segment_descriptions[0]
            jet_cmap = colormaps.get("jet")
            percents = np.linspace(0, 1, N_FRACTIONAL + 1, endpoint=True)
            colours = jet_cmap(percents)
            label = sd.SegmentLabel
            if hasattr(sd, "SegmentDescription"):
                desc = sd.SegmentDescription
            else:
                desc = None
            new_segment_descriptions = []
            new_mask_array = []
            for i in range(N_FRACTIONAL):
                new_sd = deepcopy(sd)
                p1, p2 = percents[i], percents[i + 1]
                new_sd.SegmentNumber = i + 1
                new_label = f"{label} ({int(p1*100)}-{int(p2*100)}%)"
                new_sd.SegmentLabel = new_label
                if desc is not None:
                    new_desc = f"{desc} ({int(p1*100)}-{int(p2*100)}%)"
                    new_sd.SegmentDescription = new_desc
                curr_mask = (mask_array > p1) & (mask_array <= p2)
                new_sd.RecommendedDisplayCIELabValue = rgb_to_cielab(
                    colours[i, :3] * 255
                )
                new_segment_descriptions.append(new_sd)
                new_mask_array.append(curr_mask)
            segment_descriptions = new_segment_descriptions
            mask_array = np.concatenate(new_mask_array, axis=-1).astype(bool)
            logger.info("Converted mask to pseudo-fractional DICOM seg")
        elif is_fractional:
            seg_type = hd.seg.SegmentationTypeValues.FRACTIONAL
        else:
            seg_type = hd.seg.SegmentationTypeValues.BINARY
        seg_dataset = hd.seg.Segmentation(
            source_images=image_datasets,
            pixel_array=mask_array,
            segmentation_type=seg_type,
            segment_descriptions=segment_descriptions,
            series_instance_uid=hd.UID(),
            series_number=999,
            sop_instance_uid=hd.UID(),
            instance_number=1,
            manufacturer=self.manufacturer,
            manufacturer_model_name=self.manufacturer_model_name,
            software_versions=self.algorithm_version,
            device_serial_number="42",
            series_description=seg_series_description,
        )
        seg_dataset.ClinicalTrialSeriesID = self.clinical_trial_series_id
        seg_dataset.ClinicalTrialTimePointID = self.clinical_trial_time_point_id
        seg_dataset.BodyPartExamined = self.body_part_examined

        seg_dataset.save_as(output_path)

        return "success"

    def write_dicom_rtstruct(
        self,
        mask_array: np.ndarray | sitk.Image,
        source_files: list[str],
        output_path: str,
    ):
        """
        Routine to write a DICOM RTstruct object.

        Args:
            mask_array (np.ndarray | sitk.Image): mask image.
            source_files (list[str]): list of source files.
            output_path (str): output path.

        Returns:
            str: "success" if the operation was successful.
        """
        mask_array = self.to_array_if_necessary(mask_array)
        segment_info = [
            [s.SegmentLabel, list(random_color_generator())]
            for s in self.segment_descriptions
        ]
        save_mask_as_rtstruct(
            mask_array,
            Path(source_files[0]).parent,
            output_path,
            segment_info,
        )

        return "success"

    @staticmethod
    def init_from_dcmqi_metadata_file(
        metadata_file: str,
        algorithm_version: str = "v1.0",
        manufacturer: str = "Algorithm",
        manufacturer_model_name: str = "AlgorithmModel",
    ):
        """
        Uses a DCMQI metadata file generated using [1] to initialize a
        SegWriter.

        [1] https://qiicr.org/dcmqi/#/seg

        Args:
            metadata_file (str): path to the DCMQI metadata file.
            algorithm_version (str): algorithm version.
            manufacturer (str): manufacturer.
            manufacturer_model_name (str): manufacturer model name.

        Returns:
            SegWriter: initialized SegWriter.
        """
        metadata = from_dcmqi_metainfo(metadata_file)
        segments = list(metadata.SegmentSequence)
        algorithm_name = segments[0].SegmentAlgorithmName
        algorithm_type = hd.seg.SegmentAlgorithmTypeValues[
            segments[0].SegmentAlgorithmType
        ]
        series_description = metadata.SeriesDescription
        clinical_trial_series_id = metadata.ClinicalTrialSeriesID
        clinical_trial_time_point_id = metadata.ClinicalTrialTimePointID
        body_part_examined = metadata.BodyPartExamined
        return SegWriter(
            segment_descriptions=segments,
            algorithm_name=algorithm_name,
            algorithm_version=algorithm_version,
            algorithm_family=codes.CID7162.ArtificialIntelligence,
            algorithm_type=algorithm_type,
            instance_number=1,
            series_number=999,
            series_description=series_description,
            manufacturer=manufacturer,
            manufacturer_model_name=manufacturer_model_name,
            clinical_trial_series_id=clinical_trial_series_id,
            clinical_trial_time_point_id=clinical_trial_time_point_id,
            body_part_examined=body_part_examined,
        )

    @staticmethod
    def init_from_metadata_dict(
        metadata: dict[str, str], validate: bool = False
    ):
        """
        Initializes a SegWriter from a metadata dictionary. It automatically
        detects the type of metadata and calls the appropriate initialization
        method. If there is a "path" key, it is assumed to be a DCMQI metadata
        file and the init_from_dcmqi_metadata_file method is called. Otherwise,
        the constructor is called directly.

        Args:
            metadata (dict[str, str]): metadata dictionary.
            validate (bool): whether to validate the metadata.

        Returns:
            SegWriter: initialized SegWriter.
        """
        if "path" in metadata:
            return SegWriter.init_from_dcmqi_metadata_file(metadata["path"])
        else:
            return SegWriter(**metadata, validate=validate)