longi_dataset.py

import os
from copy import deepcopy
from typing import List, Union, Type, Tuple

import numpy as np
import blosc2

from batchgenerators.utilities.file_and_folder_operations import join, load_pickle, isfile, write_pickle, subfiles, load_json
from longiseg.configuration import default_num_processes
from longiseg.training.dataloading.utils import unpack_dataset
import math

from longiseg.training.dataloading.nnunet_dataset import nnUNetBaseDataset


class LongiSegBaseDataset(nnUNetBaseDataset):
    """
    Defines the interface
    """
    def __init__(self, folder: str, identifiers: List[str] = None,
                 folder_with_segs_from_previous_stage: str = None,
                 random_prior: bool = False):
        super().__init__(folder, identifiers, folder_with_segs_from_previous_stage)
        # print('loading dataset')
        if identifiers is None:
            identifiers = self.get_identifiers(folder)
        identifiers.sort()

        patients_json = load_json(join(folder, 'patientsTr.json'))

        # this builds the subset of patients that are in the identifiers list and at least two scans
        self.patients = {}
        for patient in patients_json:
            patient_identifiers = [i for i in patients_json[patient] if i in identifiers]
            if len(patient_identifiers) > 1:
                self.patients[patient] = patient_identifiers

        self.source_folder = folder
        self.folder_with_segs_from_previous_stage = folder_with_segs_from_previous_stage
        self.identifiers = identifiers
        self.random_prior = random_prior


class LongiSegDatasetNumpy(LongiSegBaseDataset):
    def load_case(self, patient):
        current = np.random.randint(0, len(self.patients[patient]))
        if self.random_prior:
            prior = np.random.randint(0, len(self.patients[patient]))
        else:
            prior = current-1 if current > 0 else 0

        current_data_npy_file = join(self.source_folder, self.patients[patient][current] + '.npy')
        if not isfile(current_data_npy_file):
            data_current = np.load(current_data_npy_file[:-4] + '.npz')['data']
        else:
            data_current = np.load(current_data_npy_file, mmap_mode='r')

        prior_data_npy_file = join(self.source_folder, self.patients[patient][prior] + '.npy')
        if not isfile(prior_data_npy_file):
            data_prior = np.load(prior_data_npy_file[:-4] + '.npz')['data']
        else:
            data_prior = np.load(prior_data_npy_file, mmap_mode='r')

        current_seg_npy_file = join(self.source_folder, self.patients[patient][current] + '_seg.npy')
        if not isfile(current_seg_npy_file):
            seg_current = np.load(current_seg_npy_file[:-4] + '.npz')['seg']
        else:
            seg_current = np.load(current_seg_npy_file, mmap_mode='r')

        prior_seg_npy_file = join(self.source_folder, self.patients[patient][prior] + '_seg.npy')
        if not isfile(prior_seg_npy_file):
            seg_prior = np.load(prior_seg_npy_file[:-4] + '.npz')['seg']
        else:
            seg_prior = np.load(prior_seg_npy_file, mmap_mode='r')

        if self.folder_with_segs_from_previous_stage is not None:
            raise NotImplementedError("Cascade is not implemented for longitudinal segmentation")
        else:
            seg_prev = None

        properties = load_pickle(join(self.source_folder, self.patients[patient][current] + '.pkl'))
        return data_current, seg_current, data_prior, seg_prior, seg_prev, properties

    def load_single_scan(self, identifier):
        data_npy_file = join(self.source_folder, identifier + '.npy')
        if not isfile(data_npy_file):
            data = np.load(data_npy_file[:-4] + '.npz')['data']
        else:
            data = np.load(data_npy_file, mmap_mode='r')

        seg_npy_file = join(self.source_folder, identifier + '_seg.npy')
        if not isfile(seg_npy_file):
            seg = np.load(seg_npy_file[:-4] + '.npz')['seg']
        else:
            seg = np.load(seg_npy_file, mmap_mode='r')

        if self.folder_with_segs_from_previous_stage is not None:
            raise NotImplementedError("Cascade is not implemented for longitudinal segmentation")
        else:
            seg_prev = None

        properties = load_pickle(join(self.source_folder, identifier + '.pkl'))
        return data, seg, seg_prev, properties

    @staticmethod
    def save_case(
            data: np.ndarray,
            seg: np.ndarray,
            properties: dict,
            output_filename_truncated: str
    ):
        np.savez_compressed(output_filename_truncated + '.npz', data=data, seg=seg)
        write_pickle(properties, output_filename_truncated + '.pkl')

    @staticmethod
    def save_seg(
            seg: np.ndarray,
            output_filename_truncated: str
    ):
        np.savez_compressed(output_filename_truncated + '.npz', seg=seg)

    @staticmethod
    def get_identifiers(folder: str) -> List[str]:
        """
        returns all identifiers in the preprocessed data folder
        """
        case_identifiers = [i[:-4] for i in os.listdir(folder) if i.endswith("npz")]
        return case_identifiers

    @staticmethod
    def unpack_dataset(folder: str, overwrite_existing: bool = False,
                       num_processes: int = default_num_processes,
                       verify: bool = True):
        return unpack_dataset(folder, True, overwrite_existing, num_processes, verify)


class LongiSegDatasetBlosc2(LongiSegBaseDataset):
    def __init__(self, folder: str, identifiers: List[str] = None,
                 folder_with_segs_from_previous_stage: str = None,
                 random_prior: bool = False):
        super().__init__(folder, identifiers, folder_with_segs_from_previous_stage, random_prior)
        blosc2.set_nthreads(1)

    def load_case(self, patient):
        dparams = {
            'nthreads': 1
        }

        current = np.random.randint(0, len(self.patients[patient]))
        if self.random_prior:
            prior = np.random.randint(0, len(self.patients[patient]))
        else:
            prior = current-1 if current > 0 else 1

        current_data_b2nd_file = join(self.source_folder, self.patients[patient][current] + '.b2nd')
        data_current = blosc2.open(urlpath=current_data_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        prior_data_b2nd_file = join(self.source_folder, self.patients[patient][prior] + '.b2nd')
        data_prior = blosc2.open(urlpath=prior_data_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        current_seg_b2nd_file = join(self.source_folder, self.patients[patient][current] + '_seg.b2nd')
        seg_current = blosc2.open(urlpath=current_seg_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        prior_seg_b2nd_file = join(self.source_folder, self.patients[patient][prior] + '_seg.b2nd')
        seg_prior = blosc2.open(urlpath=prior_seg_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        if self.folder_with_segs_from_previous_stage is not None:
            raise NotImplementedError("Cascade is not implemented for longitudinal segmentation")
        else:
            seg_prev = None

        properties = load_pickle(join(self.source_folder, self.patients[patient][current] + '.pkl'))
        return data_current, seg_current, data_prior, seg_prior, seg_prev, properties

    def load_single_scan(self, identifier):
        dparams = {
            'nthreads': 1
        }

        data_b2nd_file = join(self.source_folder, identifier + '.b2nd')
        data = blosc2.open(urlpath=data_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        seg_b2nd_file = join(self.source_folder, identifier + '_seg.b2nd')
        seg = blosc2.open(urlpath=seg_b2nd_file, mode='r', dparams=dparams, mmap_mode='r')

        if self.folder_with_segs_from_previous_stage is not None:
            raise NotImplementedError("Cascade is not implemented for longitudinal segmentation")
        else:
            seg_prev = None

        properties = load_pickle(join(self.source_folder, identifier + '.pkl'))
        return data, seg, seg_prev, properties

    @staticmethod
    def save_case(
            data: np.ndarray,
            seg: np.ndarray,
            properties: dict,
            output_filename_truncated: str,
            chunks=None,
            blocks=None,
            chunks_seg=None,
            blocks_seg=None,
            clevel: int = 8,
            codec=blosc2.Codec.ZSTD
    ):
        blosc2.set_nthreads(1)
        if chunks_seg is None:
            chunks_seg = chunks
        if blocks_seg is None:
            blocks_seg = blocks

        cparams = {
            'codec': codec,
            # 'filters': [blosc2.Filter.SHUFFLE],
            # 'splitmode': blosc2.SplitMode.ALWAYS_SPLIT,
            'clevel': clevel,
        }
        # print(output_filename_truncated, data.shape, seg.shape, blocks, chunks, blocks_seg, chunks_seg, data.dtype, seg.dtype)
        blosc2.asarray(np.ascontiguousarray(data), urlpath=output_filename_truncated + '.b2nd', chunks=chunks,
                       blocks=blocks, cparams=cparams, mmap_mode='w+')
        blosc2.asarray(np.ascontiguousarray(seg), urlpath=output_filename_truncated + '_seg.b2nd', chunks=chunks_seg,
                       blocks=blocks_seg, cparams=cparams, mmap_mode='w+')
        write_pickle(properties, output_filename_truncated + '.pkl')

    @staticmethod
    def save_seg(
            seg: np.ndarray,
            output_filename_truncated: str,
            chunks_seg=None,
            blocks_seg=None
    ):
        blosc2.asarray(seg, urlpath=output_filename_truncated + '.b2nd', chunks=chunks_seg, blocks=blocks_seg)

    @staticmethod
    def get_identifiers(folder: str) -> List[str]:
        """
        returns all identifiers in the preprocessed data folder
        """
        case_identifiers = [i[:-5] for i in os.listdir(folder) if i.endswith(".b2nd") and not i.endswith("_seg.b2nd")]
        return case_identifiers

    @staticmethod
    def unpack_dataset(folder: str, overwrite_existing: bool = False,
                       num_processes: int = default_num_processes,
                       verify: bool = True):
        pass

    @staticmethod
    def comp_blosc2_params(
            image_size: Tuple[int, int, int, int],
            patch_size: Union[Tuple[int, int], Tuple[int, int, int]],
            bytes_per_pixel: int = 4,  # 4 byte are float32
            l1_cache_size_per_core_in_bytes=32768,  # 1 Kibibyte (KiB) = 2^10 Byte;  32 KiB = 32768 Byte
            l3_cache_size_per_core_in_bytes=1441792,
            # 1 Mibibyte (MiB) = 2^20 Byte = 1.048.576 Byte; 1.375MiB = 1441792 Byte
            safety_factor: float = 0.8  # we dont will the caches to the brim. 0.8 means we target 80% of the caches
    ):
        """
        Computes a recommended block and chunk size for saving arrays with blosc v2.

        Bloscv2 NDIM doku: "Remember that having a second partition means that we have better flexibility to fit the
        different partitions at the different CPU cache levels; typically the first partition (aka chunks) should
        be made to fit in L3 cache, whereas the second partition (aka blocks) should rather fit in L2/L1 caches
        (depending on whether compression ratio or speed is desired)."
        (https://www.blosc.org/posts/blosc2-ndim-intro/)
        -> We are not 100% sure how to optimize for that. For now we try to fit the uncompressed block in L1. This
        might spill over into L2, which is fine in our books.

        Note: this is optimized for nnU-Net dataloading where each read operation is done by one core. We cannot use threading

        Cache default values computed based on old Intel 4110 CPU with 32K L1, 128K L2 and 1408K L3 cache per core.
        We cannot optimize further for more modern CPUs with more cache as the data will need be be read by the
        old ones as well.

        Args:
            patch_size: Image size, must be 4D (c, x, y, z). For 2D images, make x=1
            patch_size: Patch size, spatial dimensions only. So (x, y) or (x, y, z)
            bytes_per_pixel: Number of bytes per element. Example: float32 -> 4 bytes
            l1_cache_size_per_core_in_bytes: The size of the L1 cache per core in Bytes.
            l3_cache_size_per_core_in_bytes: The size of the L3 cache exclusively accessible by each core. Usually the global size of the L3 cache divided by the number of cores.

        Returns:
            The recommended block and the chunk size.
        """
        # Fabians code is ugly, but eh

        num_channels = image_size[0]
        if len(patch_size) == 2:
            patch_size = [1, *patch_size]
        patch_size = np.array(patch_size)
        block_size = np.array((num_channels, *[2 ** (max(0, math.ceil(math.log2(i)))) for i in patch_size]))

        # shrink the block size until it fits in L1
        estimated_nbytes_block = np.prod(block_size) * bytes_per_pixel
        while estimated_nbytes_block > (l1_cache_size_per_core_in_bytes * safety_factor):
            # pick largest deviation from patch_size that is not 1
            axis_order = np.argsort(block_size[1:] / patch_size)[::-1]
            idx = 0
            picked_axis = axis_order[idx]
            while block_size[picked_axis + 1] == 1 or block_size[picked_axis + 1] == 1:
                idx += 1
                picked_axis = axis_order[idx]
            # now reduce that axis to the next lowest power of 2
            block_size[picked_axis + 1] = 2 ** (max(0, math.floor(math.log2(block_size[picked_axis + 1] - 1))))
            block_size[picked_axis + 1] = min(block_size[picked_axis + 1], image_size[picked_axis + 1])
            estimated_nbytes_block = np.prod(block_size) * bytes_per_pixel

        block_size = np.array([min(i, j) for i, j in zip(image_size, block_size)])

        # note: there is no use extending the chunk size to 3d when we have a 2d patch size! This would unnecessarily
        # load data into L3
        # now tile the blocks into chunks until we hit image_size or the l3 cache per core limit
        chunk_size = deepcopy(block_size)
        estimated_nbytes_chunk = np.prod(chunk_size) * bytes_per_pixel
        while estimated_nbytes_chunk < (l3_cache_size_per_core_in_bytes * safety_factor):
            if patch_size[0] == 1 and all([i == j for i, j in zip(chunk_size[2:], image_size[2:])]):
                break
            if all([i == j for i, j in zip(chunk_size, image_size)]):
                break
            # find axis that deviates from block_size the most
            axis_order = np.argsort(chunk_size[1:] / block_size[1:])
            idx = 0
            picked_axis = axis_order[idx]
            while chunk_size[picked_axis + 1] == image_size[picked_axis + 1] or patch_size[picked_axis] == 1:
                idx += 1
                picked_axis = axis_order[idx]
            chunk_size[picked_axis + 1] += block_size[picked_axis + 1]
            chunk_size[picked_axis + 1] = min(chunk_size[picked_axis + 1], image_size[picked_axis + 1])
            estimated_nbytes_chunk = np.prod(chunk_size) * bytes_per_pixel
            if np.mean([i / j for i, j in zip(chunk_size[1:], patch_size)]) > 1.5:
                # chunk size should not exceed patch size * 1.5 on average
                chunk_size[picked_axis + 1] -= block_size[picked_axis + 1]
                break
        # better safe than sorry
        chunk_size = [min(i, j) for i, j in zip(image_size, chunk_size)]

        # print(image_size, chunk_size, block_size)
        return tuple(block_size), tuple(chunk_size)


file_ending_dataset_mapping = {
    'npz': LongiSegDatasetNumpy,
    'b2nd': LongiSegDatasetBlosc2
}


def infer_dataset_class(folder: str) -> Union[Type[LongiSegDatasetBlosc2], Type[LongiSegDatasetNumpy]]:
    file_endings = set([os.path.basename(i).split('.')[-1] for i in subfiles(folder, join=False)])
    if 'pkl' in file_endings:
        file_endings.remove('pkl')
    if 'npy' in file_endings:
        file_endings.remove('npy')
    if 'json' in file_endings:
        file_endings.remove('json')
    assert len(file_endings) == 1, (f'Found more than one file ending in the folder {folder}. '
                                    f'Unable to infer nnUNetDataset variant!')
    return file_ending_dataset_mapping[list(file_endings)[0]]