U Ph4&@s.ddlmZddlZddlmZddlmZddlmZm Z ddl Z ddl Z ddl mZddlmZddlmZdd lmZmZdd lmZdd lmZdd lmZdd lmZddlmZddl m!Z!m"Z"ddl#m$Z$m%Z%m&Z&erddl'm'Z'dZ(ne&ddd\Z'Z(dgZ)e%*Z+ddZ,GdddZ-dS)) annotationsN)Callabledeepcopy) TYPE_CHECKINGAny)NdarrayOrTensor) DataLoader)Dataset)decollate_batchpad_list_data_collate)Compose)PadListDataCollate)InvertibleTransform)Invertd) Randomizable)modestack) CommonKeysPostFixoptional_import)tqdmTr)nameTestTimeAugmentationcCs|S)N)xrrV/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/test_time_augmentation.py _identity.src@sreZdZdZdedejejddeddeddf dddd d d d d d d d d d d ddZ ddZ dddddddZ dS)ra Class for performing test time augmentations. This will pass the same image through the network multiple times. The user passes transform(s) to be applied to each realization, and provided that at least one of those transforms is random, the network's output will vary. Provided that inverse transformations exist for all supplied spatial transforms, the inverse can be applied to each realization of the network's output. Once in the same spatial reference, the results can then be combined and metrics computed. Test time augmentations are a useful feature for computing network uncertainty, as well as observing the network's dependency on the applied random transforms. Reference: Wang et al., Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks, https://doi.org/10.1016/j.neucom.2019.01.103 Args: transform: transform (or composed) to be applied to each realization. At least one transform must be of type `RandomizableTrait` (i.e. `Randomizable`, `RandomizableTransform`, or `RandomizableTrait`). . All random transforms must be of type `InvertibleTransform`. batch_size: number of realizations to infer at once. num_workers: how many subprocesses to use for data. inferrer_fn: function to use to perform inference. device: device on which to perform inference. image_key: key used to extract image from input dictionary. orig_key: the key of the original input data in the dict. will get the applied transform information for this input data, then invert them for the expected data with `image_key`. orig_meta_keys: the key of the metadata of original input data, will get the `affine`, `data_shape`, etc. the metadata is a dictionary object which contains: filename, original_shape, etc. if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`. meta_key_postfix: use `key_{postfix}` to fetch the metadata according to the key data, default is `meta_dict`, the metadata is a dictionary object. For example, to handle key `image`, read/write affine matrices from the metadata `image_meta_dict` dictionary's `affine` field. this arg only works when `meta_keys=None`. to_tensor: whether to convert the inverted data into PyTorch Tensor first, default to `True`. output_device: if converted the inverted data to Tensor, move the inverted results to target device before `post_func`, default to "cpu". post_func: post processing for the inverted data, should be a callable function. return_full_data: normally, metrics are returned (mode, mean, std, vvc). Setting this flag to `True` will return the full data. Dimensions will be same size as when passing a single image through `inferrer_fn`, with a dimension appended equal in size to `num_examples` (N), i.e., `[N,C,H,W,[D]]`. progress: whether to display a progress bar. Example: .. code-block:: python model = UNet(...).to(device) transform = Compose([RandAffined(keys, ...), ...]) transform.set_random_state(seed=123) # ensure deterministic evaluation tt_aug = TestTimeAugmentation( transform, batch_size=5, num_workers=0, inferrer_fn=model, device=device ) mode, mean, std, vvc = tt_aug(test_data) rcpuTNFrintrzstr | torch.deviceboolz str | NoneNone) transform batch_size num_workers inferrer_fndevicenearest_interporig_meta_keys to_tensor output_device post_funcreturn_full_dataprogressreturnc Csb||_||_||_||_||_||_||_||_tj |_ t |j ||| | || | | d |_ | dS)N) keysr" orig_keysr(meta_key_postfixr'r)r&r+)r"r#r$r%r& image_keyr,r-rPRED _pred_keyrinverter_check_transforms)selfr"r#r$r%r&r2orig_keyr'r(r1r)r*r+r,r-rrr__init__ms* zTestTimeAugmentation.__init__cCst|jts|jgn|jj}tdd|D}tdd|D}t|dkrZtdt ||D]&\}}|rd|sdtdt |j qddS)zVShould be at least 1 random transform, and all random transforms should be invertible.cSsg|]}t|tqSr) isinstancer.0trrr sz:TestTimeAugmentation._check_transforms..cSsg|]}t|tqSr)r:rr;rrrr>srzdTTA usually has at least a `Randomizable` transform or `Compose` contains `Randomizable` transforms.zKNot all applied random transform(s) are invertible. Problematic transform: N) r:r"r transformsnparraysumwarningswarnziptype__name__)r7tsZrandomsZ invertiblesrirrrr6s z&TestTimeAugmentation._check_transforms zdict[str, Any]zQtuple[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor, float] | NdarrayOrTensor)data num_examplesr.c st||jdkrtdfddt|D}t|j}t|jjtd}g}t rjj rjt |n|D]>} |j j|j<|fddt|Dqnt|d}jr|St|dd} |d} |d} ||} | | | | fS)a Args: data: dictionary data to be processed. num_examples: number of realizations to be processed and results combined. Returns: - if `return_full_data==False`: mode, mean, std, vvc. The mode, mean and standard deviation are calculated across `num_examples` outputs at each voxel. The volume variation coefficient (VVC) is `std/mean` across the whole output, including `num_examples`. See original paper for clarification. - if `return_full_data==False`: data is returned as-is after applying the `inferrer_fn` and then concatenating across the first dimension containing `num_examples`. This allows the user to perform their own analysis if desired. rz.num_examples should be multiple of batch size.csg|] }tqSrr)r<_)drrr>sz1TestTimeAugmentation.__call__..)r$r# collate_fncs"g|]}t|jqSr)r5rinverser4)r<rJ)r7rrr>s)dim)dictr# ValueErrorranger r"r r$r has_tqdmr-rr%r2tor&r4extendr rr,rmeanstditem) r7rLrMdata_indsdloutsboutput_moderYrZZvvcr)rOr7r__call__s$     zTestTimeAugmentation.__call__)rK) rG __module__ __qualname____doc__rrIMAGELABELDEFAULT_POST_FIXr9r6rcrrrrr2s$>(*). __future__rrCcollections.abcrcopyrtypingrrnumpyr@torchmonai.config.type_definitionsrZmonai.data.dataloaderr Zmonai.data.datasetr monai.data.utilsr r monai.transforms.composer monai.transforms.croppad.batchrmonai.transforms.inverserZ monai.transforms.post.dictionaryrmonai.transforms.transformr0monai.transforms.utils_pytorch_numpy_unificationrr monai.utilsrrrrrV__all__metarirrrrrr s2