o ! i08@sddlmZddlZddlZddlZddlZddlZddlZddl Z ddl Z ddlm Z m Z ddl m Z mZddlmZddlmZddlmZddlmZdd lmZdd lmZmZmZmZddlZddl Z dd l!m"Z"dd l#m$Z$dd l%m&Z'ddl%m(Z(ddl)m*Z*ddl+m,Z,m-Z-m.Z.ddl/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5ddl6m7Z7m8Z8m9Z9m:Z:m;Z;mZ>erddl?m?Z?dZ@n eZdata_ir4r4r5 _transformZs  zDataset._transformint | slice | Sequence[int]cCsZt|tr|t|\}}}t|||}t||dSt|tjjr(t||dS| |S)z^ Returns a `Subset` if `index` is a slice or Sequence, a data item otherwise. )datasetindices) r/slicerCr:ranger collectionsabcrr@)r2r>startstopsteprCr4r4r5 __getitem__as     zDataset.__getitem__r9)r*rr+r,r-r.r-r7r>r7)r>rA)__name__ __module__ __qualname____doc__r6r<r@rKr4r4r4r5r9s    rcs.eZdZdZdfdd Zddd dZZS) DatasetFunca Execute function on the input dataset and leverage the output to act as a new Dataset. It can be used to load / fetch the basic dataset items, like the list of `image, label` paths. Or chain together to execute more complicated logic, like `partition_dataset`, `resample_datalist`, etc. The `data` arg of `Dataset` will be applied to the first arg of callable `func`. Usage example:: data_list = DatasetFunc( data="path to file", func=monai.data.load_decathlon_datalist, data_list_key="validation", base_dir="path to base dir", ) # partition dataset for every rank data_partition = DatasetFunc( data=data_list, func=lambda **kwargs: monai.data.partition_dataset(**kwargs)[torch.distributed.get_rank()], num_partitions=torch.distributed.get_world_size(), ) dataset = Dataset(data=data_partition, transform=transforms) Args: data: input data for the func to process, will apply to `func` as the first arg. func: callable function to generate dataset items. kwargs: other arguments for the `func` except for the first arg. r*rfuncrr-r.c s.tjddd||_||_||_|dS)Nr?)superr6srcrSkwargsreset)r2r*rSrV __class__r4r5r6s  zDatasetFunc.__init__N Any | NoneCallable | NonecKsJ|dur|jn|}|dur|j|fi|j|_dS||fi||_dS)aL Reset the dataset items with specified `func`. Args: data: if not None, execute `func` on it, default to `self.src`. func: if not None, execute the `func` with specified `kwargs`, default to `self.func`. kwargs: other arguments for the `func` except for the first arg. N)rUrSrVr*)r2r*rSrVrUr4r4r5rWs 8zDatasetFunc.reset)r*rrSrr-r.)NN)r*rZrSr[)rNrOrPrQr6rW __classcell__r4r4rXr5rRpsrRcsdeZdZdZededdfd'fdd Zd(ddZd)ddZddZ d d!Z d"d#Z d*d%d&Z Z S)+PersistentDataseta Persistent storage of pre-computed values to efficiently manage larger than memory dictionary format data, it can operate transforms for specific fields. Results from the non-random transform components are computed when first used, and stored in the `cache_dir` for rapid retrieval on subsequent uses. If passing slicing indices, will return a PyTorch Subset, for example: `data: Subset = dataset[1:4]`, for more details, please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.Subset The transforms which are supposed to be cached must implement the `monai.transforms.Transform` interface and should not be `Randomizable`. This dataset will cache the outcomes before the first `Randomizable` `Transform` within a `Compose` instance. For example, typical input data can be a list of dictionaries:: [{ { { 'image': 'image1.nii.gz', 'image': 'image2.nii.gz', 'image': 'image3.nii.gz', 'label': 'label1.nii.gz', 'label': 'label2.nii.gz', 'label': 'label3.nii.gz', 'extra': 123 'extra': 456 'extra': 789 }, }, }] For a composite transform like .. code-block:: python [ LoadImaged(keys=['image', 'label']), Orientationd(keys=['image', 'label'], axcodes='RAS'), ScaleIntensityRanged(keys=['image'], a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True), RandCropByPosNegLabeld(keys=['image', 'label'], label_key='label', spatial_size=(96, 96, 96), pos=1, neg=1, num_samples=4, image_key='image', image_threshold=0), ToTensord(keys=['image', 'label'])] Upon first use a filename based dataset will be processed by the transform for the [LoadImaged, Orientationd, ScaleIntensityRanged] and the resulting tensor written to the `cache_dir` before applying the remaining random dependant transforms [RandCropByPosNegLabeld, ToTensord] elements for use in the analysis. Subsequent uses of a dataset directly read pre-processed results from `cache_dir` followed by applying the random dependant parts of transform processing. During training call `set_data()` to update input data and recompute cache content. Note: The input data must be a list of file paths and will hash them as cache keys. The filenames of the cached files also try to contain the hash of the transforms. In this fashion, `PersistentDataset` should be robust to changes in transforms. This, however, is not guaranteed, so caution should be used when modifying transforms to avoid unexpected errors. If in doubt, it is advisable to clear the cache directory. Cached data is expected to be tensors, primitives, or dictionaries keying to these values. Numpy arrays will be converted to tensors, however any other object type returned by transforms will not be loadable since `torch.load` will be used with `weights_only=True` to prevent loading of potentially malicious objects. Legacy cache files may not be loadable and may need to be recomputed. Lazy Resampling: If you make use of the lazy resampling feature of `monai.transforms.Compose`, please refer to its documentation to familiarize yourself with the interaction between `PersistentDataset` and lazy resampling. pickleNTr*rr+Sequence[Callable] | Callable cache_dirPath | str | None hash_funcCallable[..., bytes] pickle_modulestrpickle_protocolr7hash_transformCallable[..., bytes] | Nonerboolr-r.c stj||d|durt|nd|_||_||_||_|jdur7|js.|jjddd|j s7t dd|_ |durC| |||_ dS)a Args: data: input data file paths to load and transform to generate dataset for model. `PersistentDataset` expects input data to be a list of serializable and hashes them as cache keys using `hash_func`. transform: transforms to execute operations on input data. cache_dir: If specified, this is the location for persistent storage of pre-computed transformed data tensors. The cache_dir is computed once, and persists on disk until explicitly removed. Different runs, programs, experiments may share a common cache dir provided that the transforms pre-processing is consistent. If `cache_dir` doesn't exist, will automatically create it. If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. pickle_module: string representing the module used for pickling metadata and objects, default to `"pickle"`. due to the pickle limitation in multi-processing of Dataloader, we can't use `pickle` as arg directly, so here we use a string name instead. if want to use other pickle module at runtime, just register like: >>> from monai.data import utils >>> utils.SUPPORTED_PICKLE_MOD["test"] = other_pickle this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save, and ``monai.data.utils.SUPPORTED_PICKLE_MOD``. pickle_protocol: specifies pickle protocol when saving, with `torch.save`. Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. hash_transform: a callable to compute hash from the transform information when caching. This may reduce errors due to transforms changing during experiments. Default to None (no hash). Other options are `pickle_hashing` and `json_hashing` functions from `monai.data.utils`. reset_ops_id: whether to set `TraceKeys.ID` to ``Tracekys.NONE``, defaults to ``True``. When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors. This is useful for skipping the transform instance checks when inverting applied operations using the cached content and with re-created transform instances. r?NTparentsexist_okzcache_dir must be a directory.)rTr6r r`rbrdrfexistsmkdiris_dirr1transform_hashset_transform_hashr) r2r*r+r`rbrdrfrgrrXr4r5r6s.     zPersistentDataset.__init__hash_xform_funcc Csg}|jjD]}t|tst|tsn||qz||}Wn(tyJ}zdt|vr2|d dd|D}||}WYd}~nd}~ww| d|_ dS)aGet hashable transforms, and then hash them. Hashable transforms are deterministic transforms that inherit from `Transform`. We stop at the first non-deterministic transform, or first that does not inherit from MONAI's `Transform` class.zis not JSON serializablermcss|]}|jjVqdSr9)rYrN).0trr4r4r5 ,z7PersistentDataset.set_transform_hash..Nutf-8) r+flatten transformsr/rrappend TypeErrorrejoindecoderq)r2rsZhashable_transformsZ_trrqtenamesr4r4r5rrs   z$PersistentDataset.set_transform_hashcCsF||_|jdur|jr!tj|jdd|jjddddSdSdS)Q Set the input data and delete all the out-dated cache content. NT) ignore_errorsrj)r*r`rnshutilrmtreeror2r*r4r4r5set_data0s zPersistentDataset.set_datacCs2|jdd}|j||dd}|jrt||S)a Process the data from original state up to the first random element. Args: item_transformed: The data to be transformed Returns: the transformed element up to the first identified random transform object cSt|tp t|t Sr9r/rrtr4r4r5Gz2PersistentDataset._pre_transform..Tend threading)r+get_index_of_firstrr2item_transformed first_randomr4r4r5_pre_transform:s z PersistentDataset._pre_transformcCs*|jdd}|dur|j||d}|S)aD Process the data from before the first random transform to the final state ready for evaluation. Args: item_transformed: The data to be transformed (already processed up to the first random transform) Returns: the transformed element through the random transforms cSrr9rrr4r4r5r[rz3PersistentDataset._post_transform..NrH)r+rrr4r4r5_post_transformOs  z!PersistentDataset._post_transformc Csd}|jdur||d}||j7}|j|d}|durt|rtztj|ddWStyE}z tj dkr;|WYd}~n3d}~wt t fys}z!dt |vsYt |t rgtd|d |n|WYd}~nd}~ww|t|}|dur|SzitZ}t||j}tjt|d d |t|jt|jd |r|sz tt ||WntyYnwWdW|SWdW|SWdW|SWdW|S1swYW|StyY|Sw) a A function to cache the expensive input data transform operations so that huge data sets (larger than computer memory) can be processed on the fly as needed, and intermediate results written to disk for future use. Args: item_transformed: The current data element to be mutated into transformed representation Returns: The transformed data_element, either from cache, or explicitly computing it. Warning: The current implementation does not encode transform information as part of the hashing mechanism used for generating cache names when `hash_transform` is None. If the transforms applied are changed in any way, the objects in the cache dir will be invalid. Nrx.ptT weights_onlywin32z"Invalid magic number; corrupt filezCorrupt cache file detected: z. Deleting and recomputing.Fconvert_numericobjfrdrf) r`rbr~rqis_filetorchloadPermissionErrorsysplatformr RuntimeErrorrer/warningswarnunlinkrrtempfileTemporaryDirectoryr namesaver r"rdrrfrmoveFileExistsError)r2rhashfile data_item_md5r3_item_transformed tmpdirnametemp_hash_filer4r4r5 _cachecheckasr               zPersistentDataset._cachecheckr>cCs||j|}||Sr9)rr*r)r2r>Zpre_random_itemr4r4r5r@s zPersistentDataset._transform)r*rr+r_r`rarbrcrdrerfr7rgrhrrir-r.)rsrcr*rrM)rNrOrPrQrrr6rrrrrrr@r\r4r4rXr5r]sA =  @r]cs>eZdZdZededdfdfdd ZddZddZZ S)CacheNTransDatasetz~ Extension of `PersistentDataset`, it can also cache the result of first N transforms, no matter it's random or not. r^NTr*rr+r_ cache_n_transr7r`rarbrcrdrerfrgrhrrir-r.c s&tj|||||||| d||_dS)a Args: data: input data file paths to load and transform to generate dataset for model. `PersistentDataset` expects input data to be a list of serializable and hashes them as cache keys using `hash_func`. transform: transforms to execute operations on input data. cache_n_trans: cache the result of first N transforms. cache_dir: If specified, this is the location for persistent storage of pre-computed transformed data tensors. The cache_dir is computed once, and persists on disk until explicitly removed. Different runs, programs, experiments may share a common cache dir provided that the transforms pre-processing is consistent. If `cache_dir` doesn't exist, will automatically create it. If `cache_dir` is `None`, there is effectively no caching. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. pickle_module: string representing the module used for pickling metadata and objects, default to `"pickle"`. due to the pickle limitation in multi-processing of Dataloader, we can't use `pickle` as arg directly, so here we use a string name instead. if want to use other pickle module at runtime, just register like: >>> from monai.data import utils >>> utils.SUPPORTED_PICKLE_MOD["test"] = other_pickle this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save, and ``monai.data.utils.SUPPORTED_PICKLE_MOD``. pickle_protocol: specifies pickle protocol when saving, with `torch.save`. Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. hash_transform: a callable to compute hash from the transform information when caching. This may reduce errors due to transforms changing during experiments. Default to None (no hash). Other options are `pickle_hashing` and `json_hashing` functions from `monai.data.utils`. reset_ops_id: whether to set `TraceKeys.ID` to ``Tracekys.NONE``, defaults to ``True``. When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors. This is useful for skipping the transform instance checks when inverting applied operations using the cached content and with re-created transform instances. )r*r+r`rbrdrfrgrN)rTr6r) r2r*r+rr`rbrdrfrgrrXr4r5r6s0 zCacheNTransDataset.__init__cCs|j||jdd}t||S)z Process the data from original state up to the N element. Args: item_transformed: The data to be transformed Returns: the transformed element up to the N transform object Tr)r+rrr2rr4r4r5rs z!CacheNTransDataset._pre_transformcCs|j||jdS)a Process the data from before the N + 1 transform to the final state ready for evaluation. Args: item_transformed: The data to be transformed (already processed up to the first N transform) Returns: the final transformed result r)r+rrr4r4r5rs z"CacheNTransDataset._post_transform)r*rr+r_rr7r`rarbrcrdrerfr7rgrhrrir-r.) rNrOrPrQrrr6rrr\r4r4rXr5rs <rcspeZdZdZdeddedddfd'fdd Zd(fdd ZddZdd Z d)d!d"Z fd#d$Z d%d&Z Z S)* LMDBDataseta Extension of `PersistentDataset` using LMDB as the backend. See Also: :py:class:`monai.data.PersistentDataset` Examples: >>> items = [{"data": i} for i in range(5)] # [{'data': 0}, {'data': 1}, {'data': 2}, {'data': 3}, {'data': 4}] >>> lmdb_ds = monai.data.LMDBDataset(items, transform=monai.transforms.SimulateDelayd("data", delay_time=1)) >>> print(list(lmdb_ds)) # using the cached results cacheZ monai_cacheTNr*rr+r_r` Path | strrbrcdb_namereprogressrirgrhr lmdb_kwargs dict | Noner-r.c stj||||||| d||_|jstd|j|d|_| p#i|_|jdds1d|jd<d|_|j |jdt d |j d dS) a Args: data: input data file paths to load and transform to generate dataset for model. `LMDBDataset` expects input data to be a list of serializable and hashes them as cache keys using `hash_func`. transform: transforms to execute operations on input data. cache_dir: if specified, this is the location for persistent storage of pre-computed transformed data tensors. The cache_dir is computed once, and persists on disk until explicitly removed. Different runs, programs, experiments may share a common cache dir provided that the transforms pre-processing is consistent. If the cache_dir doesn't exist, will automatically create it. Defaults to "./cache". hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. db_name: lmdb database file name. Defaults to "monai_cache". progress: whether to display a progress bar. pickle_protocol: specifies pickle protocol when saving, with `torch.save`. Defaults to torch.serialization.DEFAULT_PROTOCOL. For more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. hash_transform: a callable to compute hash from the transform information when caching. This may reduce errors due to transforms changing during experiments. Default to None (no hash). Other options are `pickle_hashing` and `json_hashing` functions from `monai.data.utils`. reset_ops_id: whether to set `TraceKeys.ID` to ``Tracekeys.NONE``, defaults to ``True``. When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors. This is useful for skipping the transform instance checks when inverting applied operations using the cached content and with re-created transform instances. lmdb_kwargs: additional keyword arguments to the lmdb environment. for more details please visit: https://lmdb.readthedocs.io/en/release/#environment-class )r*r+r`rbrfrgrzcache_dir must be specified.z.lmdbmap_sizerlN show_progresszAccessing lmdb file: .) rTr6rr`r1db_filerget _read_env_fill_cache_start_readerprintabsolute) r2r*r+r`rbrrrfrgrrrXr4r5r6s&)   zLMDBDataset.__init__cs"tj|d|j|jd|_dS)r)r*rN)rTrrrrrrXr4r5rTszLMDBDataset.set_datacCs.t}tjt|||jd|d|S)N)rfr)rrrr rfseekread)r2valoutr4r4r5_safe_serialize\s zLMDBDataset._safe_serializecCstjt|dddS)NcpuT) map_locationr)rrr)r2rr4r4r5_safe_deserializebzLMDBDataset._safe_deserializec Csvd|jd<tjd|jdd|j}|rtstd|jdd}tr,|r,t|j n|j D]}| |}d\}}}|s|dkrzM| } | |}Wdn1sVwY|r_Wq;|duro| t|}||}|jd d} | ||Wdn1swYd }WnBtjyd|d }}|d } | d } td t| d?dt| d?d|| Yntjy|dYnw|s|dksA|s|d } |td| dq/Wdn1swY|d } |d |jd<| |jd <|jdddurd|jd<|jdddur-d|jd<tjd|jdd|jS)aF Check the LMDB cache and write the cache if needed. py-lmdb doesn't have a good support for concurrent write. This method can be used with multiple processes, but it may have a negative impact on the performance. Args: show_progress: whether to show the progress bar if possible. Freadonly)pathsubdirzHLMDBDataset: tqdm is not installed. not displaying the caching progress.write)FNrNTrz!Resizing the cache database from zMB to zMB.z;LMDB map size reached, increase size above current size of rlock readaheadr4)rr(openrhas_tqdmrrbeginr&r*rbcursorZset_keyrrrputZ MapFullErrorinfor7Z set_mapsizeZMapResizedErrorcloser1r) r2renvZ search_txnitemkeydoneretryrrtxnsizenew_sizer4r4r5resd             !    z$LMDBDataset._fill_cache_start_readerc s|jdur |jdd|_|jjdd}|||}Wdn1s&wY|dur:tdt|Sz| |WSt yQ}zt d|d}~ww)zq if the item is not found in the lmdb file, resolves to the persistent cache default behaviour. NFrrz;LMDBDataset: cache key not found, running fallback caching.z)Invalid cache value, corrupted lmdb file?) rrrrrbrrrTrrr0r)r2rrr*errrXr4r5rs     zLMDBDataset._cachecheckcCsD|jdur ||_t|j}t|j|d<|j|d<|S)z4 Returns: dataset info dictionary. Nrfilename)rrdictrr:r*rr)r2rr4r4r5rs  zLMDBDataset.info)r*rr+r_r`rrbrcrrerrirgrhrrirrr-r.r)T)rNrOrPrQrrr6rrrrrrr\r4r4rXr5rs"@  9rc sheZdZdZdejddddddedf d)fdd Zd*ddZd+d,d!d"Z d-d$d%Z d.fd'd( Z Z S)/ CacheDatasetau Dataset with cache mechanism that can load data and cache deterministic transforms' result during training. By caching the results of non-random preprocessing transforms, it accelerates the training data pipeline. If the requested data is not in the cache, all transforms will run normally (see also :py:class:`monai.data.dataset.Dataset`). Users can set the cache rate or number of items to cache. It is recommended to experiment with different `cache_num` or `cache_rate` to identify the best training speed. The transforms which are supposed to be cached must implement the `monai.transforms.Transform` interface and should not be `Randomizable`. This dataset will cache the outcomes before the first `Randomizable` `Transform` within a `Compose` instance. So to improve the caching efficiency, please always put as many as possible non-random transforms before the randomized ones when composing the chain of transforms. If passing slicing indices, will return a PyTorch Subset, for example: `data: Subset = dataset[1:4]`, for more details, please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.Subset For example, if the transform is a `Compose` of:: transforms = Compose([ LoadImaged(), EnsureChannelFirstd(), Spacingd(), Orientationd(), ScaleIntensityRanged(), RandCropByPosNegLabeld(), ToTensord() ]) when `transforms` is used in a multi-epoch training pipeline, before the first training epoch, this dataset will cache the results up to ``ScaleIntensityRanged``, as all non-random transforms `LoadImaged`, `EnsureChannelFirstd`, `Spacingd`, `Orientationd`, `ScaleIntensityRanged` can be cached. During training, the dataset will load the cached results and run ``RandCropByPosNegLabeld`` and ``ToTensord``, as ``RandCropByPosNegLabeld`` is a randomized transform and the outcome not cached. During training call `set_data()` to update input data and recompute cache content, note that it requires `persistent_workers=False` in the PyTorch DataLoader. Note: `CacheDataset` executes non-random transforms and prepares cache content in the main process before the first epoch, then all the subprocesses of DataLoader will read the same cache content in the main process during training. it may take a long time to prepare cache content according to the size of expected cache data. So to debug or verify the program before real training, users can set `cache_rate=0.0` or `cache_num=0` to temporarily skip caching. Lazy Resampling: If you make use of the lazy resampling feature of `monai.transforms.Compose`, please refer to its documentation to familiarize yourself with the interaction between `CacheDataset` and lazy resampling. N?rTFr*rr+r, cache_numr7 cache_ratefloat num_workers int | Nonerri copy_cache as_contiguous hash_as_keyrbrc runtime_cachebool | str | list | ListProxyr-r.c stj||d||_||_||_||_||_| |_| |_||_ |j dur.t t |j d|_ | |_ d|_ g|_g|_||dS)a Args: data: input data to load and transform to generate dataset for model. transform: transforms to execute operations on input data. cache_num: number of items to be cached. Default is `sys.maxsize`. will take the minimum of (cache_num, data_length x cache_rate, data_length). cache_rate: percentage of cached data in total, default is 1.0 (cache all). will take the minimum of (cache_num, data_length x cache_rate, data_length). num_workers: the number of worker threads if computing cache in the initialization. If num_workers is None then the number returned by os.cpu_count() is used. If a value less than 1 is specified, 1 will be used instead. progress: whether to display a progress bar. copy_cache: whether to `deepcopy` the cache content before applying the random transforms, default to `True`. if the random transforms don't modify the cached content (for example, randomly crop from the cached image and deepcopy the crop region) or if every cache item is only used once in a `multi-processing` environment, may set `copy=False` for better performance. as_contiguous: whether to convert the cached NumPy array or PyTorch tensor to be contiguous. it may help improve the performance of following logic. hash_as_key: whether to compute hash value of input data as the key to save cache, if key exists, avoid saving duplicated content. it can help save memory when the dataset has duplicated items or augmented dataset. hash_func: if `hash_as_key`, a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. runtime_cache: mode of cache at the runtime. Default to `False` to prepare the cache content for the entire ``data`` during initialization, this potentially largely increase the time required between the constructor called and first mini-batch generated. Three options are provided to compute the cache on the fly after the dataset initialization: 1. ``"threads"`` or ``True``: use a regular ``list`` to store the cache items. 2. ``"processes"``: use a ListProxy to store the cache items, it can be shared among processes. 3. A list-like object: a users-provided container to be used to store the cache items. For `thread-based` caching (typically for caching cuda tensors), option 1 is recommended. For single process workflows with multiprocessing data loading, option 2 is recommended. For multiprocessing workflows (typically for distributed training), where this class is initialized in subprocesses, option 3 is recommended, and the list-like object should be prepared in the main process and passed to all subprocesses. Not following these recommendations may lead to runtime errors or duplicated cache across processes. r?Nrr)rTr6set_numset_raterrrrrbrmaxr7rr_cache _hash_keysr) r2r*r+rrrrrrrrbrrXr4r5r6s 7 zCacheDataset.__init__cs|_d fdd }jr5fddtjD}|t|t|dj_t|dj}n|tjttj}j dvrP |_ dSt j t rhd j vrhtdgj_ dSj d usxt j t rd j vrdgj_ dSj _ dS) aA Set the input data and run deterministic transforms to generate cache content. Note: should call this func after an entire epoch and must set `persistent_workers=False` in PyTorch DataLoader, because it needs to create new worker processes based on new generated cache content. data_lenr7cs"ttjt|j|_dSr9)minr7rrr)rr;r4r5_compute_cache_numHs"z1CacheDataset.set_data.._compute_cache_numcsi|] \}}||qSr4)rb)rtivr;r4r5 Msz)CacheDataset.set_data..N)FNprocessTthread)rr7)r*r enumerater:listrrvaluesrEr _fill_cacherr/rer)r2r*rmappingrCr4r;r5r=s(     zCacheDataset.set_datar cCs|jdkrgS|durtt|j}|jrtstdt|j-}|jr>tr>tt | |j |t |ddWdSt| |j |WdS1sQwYdS)z Compute and fill the cache content from data source. Args: indices: target indices in the `self.data` source to compute cache. if None, use the first `cache_num` items. rNz>tqdm is not installed, will not show the caching progress bar.zLoading dataset)totaldesc) rr rErrrrr rr&imap_load_cache_itemr:)r2rCpr4r4r5r bs     $zCacheDataset._fill_cacheidxcCsB|j|}|jdd}|j||dd}|jrt|tjd}|S)zN Args: idx: the index of the input data sequence. cSrr9rrr4r4r5r~rz/CacheDataset._load_cache_item..Tr) memory_format)r*r+rrrrcontiguous_format)r2rrrr4r4r5rvs zCacheDataset._load_cache_itemr>csd}|jr||j|}||jvr|j|}n |t||jkr$|}|dur.t|S|j dur7t d|j |}|durJ| |}|j |<t |j tsTtd|j dd}|durr|jdurit|n|}|j ||d}|S)Nz@cache buffer is not initialized, please call `set_data()` first.z:transform must be an instance of monai.transforms.Compose.cSrr9rrr4r4r5rrz)CacheDataset._transform..Tr)rrbr*rr>r:rrTr@rrrr/r+rr1rrr)r2r>Z cache_indexrr*rrXr4r5r@s0      zCacheDataset._transform)r*rr+r,rr7rrrrrrirrirrirrirbrcrrr-r.r*rr-r.r9)r-r )rr7rM) rNrOrPrQrmaxsizerr6rr rr@r\r4r4rXr5rs"9 H % rc seZdZdZddejdddddddddf d>fdd Zd?fdd Zd@d!d"ZdAd#d$Z d%d&Z d'd(Z d)d*Z d+d,Z d-d.Zd/d0Zd1d2ZdBd4d5Zd6d7Zd8d9ZdAd:d;Zd(z.SmartCacheDataset.__init__..) set_random_stater _start_posrLock _update_lock_round _replace_done _replace_mgrNotImplementedErrorrTr6rr rr:rrr1rrr7 _total_numrmathceil _replace_numrE _replacementsr _replace_data_idx_compute_data_idx)r2r*r+rrrrrrrrrrrrXr4r5r6sJ     " zSmartCacheDataset.__init__csB|r td||jrt|}||t|dS)z Set the input data and run deterministic transforms to generate cache content. Note: should call `shutdown()` before calling this func. zcCs|j|}|||j|<dS)zT Execute deterministic transforms on the new data for replacement. N)r1rr0)r2r>r7r4r4r5_replace_cache_threads z'SmartCacheDataset._replace_cache_threadcCsLt|j}||jtt|jWdn1swYd|_dS)z Compute expected items for the replacement of next epoch, execute deterministic transforms. It can support multi-threads to accelerate the computation progress. NT)r rmaprEr rEr/r))r2rr4r4r5_compute_replacementss  z'SmartCacheDataset._compute_replacementscCsh|j'|jdkrd|_ WddS|j|kr|d|jfWdS1s-wYdS)zX Wait thread lock and replace training items in the background thread. rTN)TF)r'r(r)rG)r2 check_roundr4r4r5_try_manage_replacements  $z)SmartCacheDataset._try_manage_replacementcCs0d}d}|s||\}}td|rdSdS)z5 Background thread for replacement. rHFr@N)rJrArB)r2rIrr4r4r5r=s   z$SmartCacheDataset.manage_replacementcC|jS)zQ The dataset length is given by cache_num instead of len(data). )rr;r4r4r5r<szSmartCacheDataset.__len__)r*rr+r,rrrr7rrrrrrrrirrirr7rrirrir-r.rr)r-r.rM)rNrOrPrQrrr6rr5r2r3rHr9r?rCrDr4rErGrJr=r<r\r4r4rXr5rs<K:        rcs8eZdZdZddfd d Zdd d ZdddZZS) ZipDatasetaS Zip several PyTorch datasets and output data(with the same index) together in a tuple. If the output of single dataset is already a tuple, flatten it and extend to the result. For example: if datasetA returns (img, imgmeta), datasetB returns (seg, segmeta), finally return (img, imgmeta, seg, segmeta). And if the datasets don't have same length, use the minimum length of them as the length of ZipDataset. If passing slicing indices, will return a PyTorch Subset, for example: `data: Subset = dataset[1:4]`, for more details, please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.Subset Examples:: >>> zip_data = ZipDataset([[1, 2, 3], [4, 5]]) >>> print(len(zip_data)) 2 >>> for item in zip_data: >>> print(item) [1, 4] [2, 5] Ndatasetsrr+r[r-r.cstjt||ddS)z Args: datasets: list of datasets to zip together. transform: a callable data transform operates on the zipped item from `datasets`. )r+N)rTr6r )r2rMr+rXr4r5r6szZipDataset.__init__r7cCstdd|jDS)Ncss|]}t|VqdSr9)r:)rtrBr4r4r5rvrwz%ZipDataset.__len__..)rr*r;r4r4r5r<rzZipDataset.__len__r>cCsNdd}g}|jD] }||||q |jdur#d|j_||}t|S)NcSst|ttfr t|S|gSr9)r/tupler )xr4r4r5to_listsz&ZipDataset._transform..to_listF)r*r>r+Z map_itemsrN)r2r>rPr*rBr4r4r5r@s   zZipDataset._transformr9)rMrr+r[r-r.rLrMrNrOrPrQr6r<r@r\r4r4rXr5rLs  rLc@sFeZdZdZ     ddddZdddZddddZdddZdS) ArrayDataseta4 Dataset for segmentation and classification tasks based on array format input data and transforms. It ensures the same random seeds in the randomized transforms defined for image, segmentation and label. The `transform` can be :py:class:`monai.transforms.Compose` or any other callable object. For example: If train based on Nifti format images without metadata, all transforms can be composed:: img_transform = Compose( [ LoadImage(image_only=True), EnsureChannelFirst(), RandAdjustContrast() ] ) ArrayDataset(img_file_list, img_transform=img_transform) If training based on images and the metadata, the array transforms can not be composed because several transforms receives multiple parameters or return multiple values. Then Users need to define their own callable method to parse metadata from `LoadImage` or set `affine` matrix to `Spacing` transform:: class TestCompose(Compose): def __call__(self, input_): img, metadata = self.transforms[0](input_) img = self.transforms[1](img) img, _, _ = self.transforms[2](img, metadata["affine"]) return self.transforms[3](img), metadata img_transform = TestCompose( [ LoadImage(image_only=False), EnsureChannelFirst(), Spacing(pixdim=(1.5, 1.5, 3.0)), RandAdjustContrast() ] ) ArrayDataset(img_file_list, img_transform=img_transform) Examples:: >>> ds = ArrayDataset([1, 2, 3, 4], lambda x: x + 0.1) >>> print(ds[0]) 1.1 >>> ds = ArrayDataset(img=[1, 2, 3, 4], seg=[5, 6, 7, 8]) >>> print(ds[0]) [1, 5] Nimgr img_transformr[segSequence | None seg_transformlabelslabel_transformr-r.c CsZ||f||f||fg}|jtddd|D}t|dkr#|dnt||_d|_dS)a Initializes the dataset with the filename lists. The transform `img_transform` is applied to the images and `seg_transform` to the segmentations. Args: img: sequence of images. img_transform: transform to apply to each element in `img`. seg: sequence of segmentations. seg_transform: transform to apply to each element in `seg`. labels: sequence of labels. label_transform: transform to apply to each element in `labels`. rcSs*g|]}|ddurt|d|dqS)rNrr)rtrOr4r4r5r"Us*z)ArrayDataset.__init__..rrN)r$r!r:rLrB_seed) r2rSrTrUrWrXrYitemsrMr4r4r5r6=s  zArrayDataset.__init__r7cCr8r9)r:rBr;r4r4r5r<Zr=zArrayDataset.__len__r*rZcCst|jjtdd|_dS)Nuint32)dtype)r7r6randintrrZrr4r4r5r5]szArrayDataset.randomizer>cCsv|t|jtr#|jjD]}t|dd}t|tr"|j|jdqt|jdd}t|tr6|j|jd|j|S)Nr+r) r5r/rBrLr*getattrrr$rZ)r2r>rBr+r4r4r5rK`s      zArrayDataset.__getitem__)NNNNN)rSrrTr[rUrVrWr[rXrVrYr[r-r.rLr9)r*rZr-r.rM)rNrOrPrQr6r<r5rKr4r4r4r5rR s4  rRcs:eZdZdZ  ddfd d ZddZdddZZS)NPZDictItemDataseta= Represents a dataset from a loaded NPZ file. The members of the file to load are named in the keys of `keys` and stored under the keyed name. All loaded arrays must have the same 0-dimension (batch) size. Items are always dicts mapping names to an item extracted from the loaded arrays. If passing slicing indices, will return a PyTorch Subset, for example: `data: Subset = dataset[1:4]`, for more details, please check: https://pytorch.org/docs/stable/data.html#torch.utils.data.Subset Args: npzfile: Path to .npz file or stream containing .npz file data keys: Maps keys to load from file to name to store in dataset transform: Transform to apply to batch dict other_keys: secondary data to load from file and store in dict `other_keys`, not returned by __getitem__ Nr4npzfilestr | IOkeysdict[str, str]r+$Callable[..., dict[str, Any]] | None other_keysSequence[str] | Nonec st|tr|nd|_t||_t|fdd|jD|_|jt tt |j j d|_ |dur8infdd|D|_|jD]\}}|j d|j krdtd|j d|d|j dqGtg|dS) NZSTREAMcsi|] \}}||qSr4r4)rtZdatakZstoredkdatr4r5rz/NPZDictItemDataset.__init__..rcsi|]}||qSr4r4)rtkrhr4r5rrz:All loaded arrays must have the same first dimension size z , array `z ` has size )r/rerarrcnprr[arraysrr%r shapelengthrfr1rTr6)r2rarcr+rfrkrrXrhr5r6}s&  " zNPZDictItemDataset.__init__cCrKr9)ror;r4r4r5r<szNPZDictItemDataset.__len__r>r7cs^fdd|jD}|jdur||n|}t|ts)t|tr+t|dtr+|Std)Ncsi|] \}}||qSr4r4)rtrkrr>r4r5rrjz1NPZDictItemDataset._transform..rzIWith a dict supplied to Compose, should return a dict or a list of dicts.)rmr[r+r/rr AssertionError)r2r>r*resultr4rpr5r@s "zNPZDictItemDataset._transform)Nr4)rarbrcrdr+rerfrgrMrQr4r4rXr5r`nsr`cs2eZdZdZ       ddfdd ZZS) CSVDataseta Dataset to load data from CSV files and generate a list of dictionaries, every dictionary maps to a row of the CSV file, and the keys of dictionary map to the column names of the CSV file. It can load multiple CSV files and join the tables with additional `kwargs` arg. Support to only load specific rows and columns. And it can also group several loaded columns to generate a new column, for example, set `col_groups={"meta": ["meta_0", "meta_1", "meta_2"]}`, output can be:: [ {"image": "./image0.nii", "meta_0": 11, "meta_1": 12, "meta_2": 13, "meta": [11, 12, 13]}, {"image": "./image1.nii", "meta_0": 21, "meta_1": 22, "meta_2": 23, "meta": [21, 22, 23]}, ] Args: src: if provided the filename of CSV file, it can be a str, URL, path object or file-like object to load. also support to provide pandas `DataFrame` directly, will skip loading from filename. if provided a list of filenames or pandas `DataFrame`, it will join the tables. row_indices: indices of the expected rows to load. it should be a list, every item can be a int number or a range `[start, end)` for the indices. for example: `row_indices=[[0, 100], 200, 201, 202, 300]`. if None, load all the rows in the file. col_names: names of the expected columns to load. if None, load all the columns. col_types: `type` and `default value` to convert the loaded columns, if None, use original data. it should be a dictionary, every item maps to an expected column, the `key` is the column name and the `value` is None or a dictionary to define the default value and data type. the supported keys in dictionary are: ["type", "default"]. for example:: col_types = { "subject_id": {"type": str}, "label": {"type": int, "default": 0}, "ehr_0": {"type": float, "default": 0.0}, "ehr_1": {"type": float, "default": 0.0}, "image": {"type": str, "default": None}, } col_groups: args to group the loaded columns to generate a new column, it should be a dictionary, every item maps to a group, the `key` will be the new column name, the `value` is the names of columns to combine. for example: `col_groups={"ehr": [f"ehr_{i}" for i in range(10)], "meta": ["meta_1", "meta_2"]}` transform: transform to apply on the loaded items of a dictionary data. kwargs_read_csv: dictionary args to pass to pandas `read_csv` function. kwargs: additional arguments for `pandas.merge()` API to join tables. NrUstr | Sequence[str] | None row_indicesSequence[int | str] | None col_namesrg col_types'dict[str, dict[str, Any] | None] | None col_groupsdict[str, Sequence[str]] | Noner+r[kwargs_read_csvrc  st|ttfs |fn|} g} | D]*} t| tr+| |r$tj| fi|nt| qt| tjr7| | qtdt d| ||||d|} t j | |ddS)Nz.`src` must be file path or pandas `DataFrame`.)dfsrurwrxrzr?r4) r/rNr rer{pdread_csv DataFramer1rrTr6) r2rUrurwrxrzr+r|rVZsrcsr}rr*rXr4r5r6s  (   zCSVDataset.__init__)NNNNNNN)rUrtrurvrwrgrxryrzr{r+r[r|r)rNrOrPrQr6r\r4r4rXr5rss1rscsBeZdZdZeddfdfdd ZddZddZddZZ S) GDSDataseta An extension of the PersistentDataset using direct memory access(DMA) data path between GPU memory and storage, thus avoiding a bounce buffer through the CPU. This direct path can increase system bandwidth while decreasing latency and utilization load on the CPU and GPU. A tutorial is available: https://github.com/Project-MONAI/tutorials/blob/main/modules/GDS_dataset.ipynb. See also: https://github.com/rapidsai/kvikio NTr*rr+r_r`radevicer7rbrcrgrhrrirVrr-r.c s0tjd||||||d|||_i|_dS)aM Args: data: input data file paths to load and transform to generate dataset for model. `GDSDataset` expects input data to be a list of serializable and hashes them as cache keys using `hash_func`. transform: transforms to execute operations on input data. cache_dir: If specified, this is the location for gpu direct storage of pre-computed transformed data tensors. The cache_dir is computed once, and persists on disk until explicitly removed. Different runs, programs, experiments may share a common cache dir provided that the transforms pre-processing is consistent. If `cache_dir` doesn't exist, will automatically create it. If `cache_dir` is `None`, there is effectively no caching. device: target device to put the output Tensor data. Note that only int can be used to specify the gpu to be used. hash_func: a callable to compute hash from data items to be cached. defaults to `monai.data.utils.pickle_hashing`. hash_transform: a callable to compute hash from the transform information when caching. This may reduce errors due to transforms changing during experiments. Default to None (no hash). Other options are `pickle_hashing` and `json_hashing` functions from `monai.data.utils`. reset_ops_id: whether to set `TraceKeys.ID` to ``Tracekys.NONE``, defaults to ``True``. When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors. This is useful for skipping the transform instance checks when inverting applied operations using the cached content and with re-created transform instances. )r*r+r`rbrgrNr4)rTr6r _meta_cache) r2r*r+r`rrbrgrrVrXr4r5r6s$ zGDSDataset.__init__c Csd}|jdur||d}||j7}|j|d}|dur@|r@tj|j t |t ri}|D]=}|j |j d|dd}t j|d||dtdd ||<t|||d d |jd ||<|||d <q8|WdSt |tjtjfr|j |j dd}t j||dtdd }t||d d |jd }ttdd|} t| r||fWdS|WdSddtt|D}t|D]L\} } | D]E}|j |j d|d| d} t j|d|d| | dtdd } t|| | d d |jd } || || |d | iqq|WdS1s;wY|t|}|durN|St |t r|D],}|d|}|j d|d}t ||tjtjfr| ||||qV|SnGt |tjtjfr|}|j d}| |||n,t|D]'\} } | D]}|d|d| }|j d|d| }| | ||qqt!|d"|S)a In order to enable direct storage to the GPU when loading the hashfile, rewritten this function. Note that in this function, it will always return `torch.Tensor` when load data from cache. Args: item_transformed: The current data element to be mutated into transformed representation Returns: The transformed data_element, either from cache, or explicitly computing it. Warning: The current implementation does not encode transform information as part of the hashing mechanism used for generating cache names when `hash_transform` is None. If the transforms applied are changed in any way, the objects in the cache dir will be invalid. Nrxr-z-meta)meta_hash_file_namer]r4)r]likernzcuda:)rZ _meta_dictcSs|dvS)N)r]rnr4)rr4r4r5rNsz(GDSDataset._cachecheck..cSsg|]}iqSr4r4r r4r4r5r"Sr#z*GDSDataset._cachecheck..z-meta-a)#r`rbr~rqrcpcudaDevicerr/r_load_meta_cacher kvikio_numpyfromfileemptyr reshaperlndarrayrTensorr filterrcrirEr:r updaterr_create_new_cacherr)r2rrrrrkZmeta_k_meta_dataZ filtered_keysrZ_itemZmeta_i_kZitem_kr data_hashfilerr4r4r5r(sz   (&  ""   zGDSDataset._cachecheckc Cslt|tr t|jni|j|<t|tr|jn|}t|tjr#|}|j |j|d<t |j |j|d<t ||zpta}|j|}t||}tjt|j|dd|t|jt|jd|r}|sz tt ||Wnty|YnwWdWdSWdWdSWdWdSWdWdS1swYWdStyYdSw)Nrnr]Frr)r/rrmetararrayrrnumpyrnrer]rtofilerrr`r rr r"rdrrfrrrrr)r2r*rrZ_item_transformed_datarZmeta_hash_filerr4r4r5rwsD        & zGDSDataset._create_new_cachecCs(||jvr |j|Stj|j|ddS)NTr)rrrr`)r2rr4r4r5rs  zGDSDataset._load_meta_cache)r*rr+r_r`rarr7rbrcrgrhrrirVrr-r.) rNrOrPrQrr6rrrr\r4r4rXr5rs0O r)Q __future__rcollections.abcrFr-rrrrrArrrrriorZmultiprocessing.managersrZmultiprocessing.poolr pathlibr r^r typingr r rrrrlrtorch.multiprocessingrZtorch.serializationrtorch.utils.datarZ _TorchDatasetrmonai.data.meta_tensorrmonai.data.utilsrrrmonai.transformsrrrrrr monai.utilsrr r!r"r#r$monai.utils.miscr%r&rrr!r(r~rrRr]rrrrrLrRr`rsrr4r4r4r5sl                   72^;l22c4K