U PhVk@sddlmZddlZddlZddlmZddlmZmZddl Z ddl Z ddl Z ddl mZddlmZmZddlmZmZmZmZddlmZdd lmZmZmZmZdd lmZm Z m!Z!m"Z"d gZ#e$dd d Z%ddZ&Gdd d ee j'Z(dS)) annotationsN)deepcopy)AnySequence) NdarrayTensor)MetaObjget_track_meta)affine_to_spacingdecollate_batchlist_data_collateremove_extra_metadata)look_up_option)LazyAttrMetaKeysPostFix SpaceKeys)convert_data_typeconvert_to_dst_typeconvert_to_numpyconvert_to_tensor MetaTensorcCsHttdrDt|drDttj|jrDtttj|jtrDttj|jSdS)N return_types__name__)hasattrtorchrr isinstancegetattrtype)funcrK/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/meta_tensor.py_get_named_tuple_like_type!s r!c CsHt|ttttjtjtjtj fpFt|t pDt|t oDt dd|D S)Ncss|]}t|tVqdS)N)rr.0xrrr /sz*_not_requiring_metadata..) rintstrbytesrSizedtypedevicenpndarrayrrany)retrrr _not_requiring_metadata-s $r0cseZdZdZedQdddddddZdRdddd dfd d Zed d d ddZeddZ edSddfdd Z eddZ ddZ ddZ eejfddddZddd d!Zejddfd"d#ZdTd%d&d'd(Zed)d*Zejd dd+d*Zejdfd,d-d.d/d0ZdUd1d2Zeddd3d4Zejd5d d6d7d4Zed8d9Zd:d;Zdd?Z dVd@dAZ!dBdCZ"edWd5dd%dEd,dFdGdHZ#dIdJZ$dKdLZ%dMdNZ&d ddOdPZ'Z(S)Xra Class that inherits from both `torch.Tensor` and `MetaObj`, adding support for metadata. Metadata is stored in the form of a dictionary. Nested, an affine matrix will be stored. This should be in the form of `torch.Tensor`. Behavior should be the same as `torch.Tensor` aside from the extended meta functionality. Copying of information: * For `c = a + b`, then auxiliary data (e.g., metadata) will be copied from the first instance of `MetaTensor` if `a.is_batch` is False (For batched data, the metadata will be shallow copied for efficiency purposes). Example: .. code-block:: python import torch from monai.data import MetaTensor t = torch.tensor([1,2,3]) affine = torch.as_tensor([[2,0,0,0], [0,2,0,0], [0,0,2,0], [0,0,0,1]], dtype=torch.float64) meta = {"some": "info"} m = MetaTensor(t, affine=affine, meta=meta) m2 = m + m assert isinstance(m2, MetaTensor) assert m2.meta["some"] == "info" assert torch.all(m2.affine == affine) Notes: - Requires pytorch 1.9 or newer for full compatibility. - Older versions of pytorch (<=1.8), `torch.jit.trace(net, im)` may not work if `im` is of type `MetaTensor`. This can be resolved with `torch.jit.trace(net, im.as_tensor())`. - For pytorch < 1.8, sharing `MetaTensor` instances across processes may not be supported. - For pytorch < 1.9, next(iter(meta_tensor)) returns a torch.Tensor. see: https://github.com/pytorch/pytorch/issues/54457 - A warning will be raised if in the constructor `affine` is not `None` and `meta` already contains the key `affine`. - You can query whether the `MetaTensor` is a batch with the `is_batch` attribute. - With a batch of data, `batch[0]` will return the 0th image with the 0th metadata. When the batch dimension is non-singleton, e.g., `batch[:, 0]`, `batch[..., -1]` and `batch[1:3]`, then all (or a subset in the last example) of the metadata will be returned, and `is_batch` will return `True`. - When creating a batch with this class, use `monai.data.DataLoader` as opposed to `torch.utils.data.DataLoader`, as this will take care of collating the metadata properly. Nztorch.Tensor | Nonez dict | Nonez list | None)affinemetaapplied_operationsreturncOs:|r|dd|dddni}tj|f|||S)Nr+r*r+r*)popr as_tensor as_subclass)clsr$r1r2r3argskwargs_kwargsrrr __new__is "zMetaTensor.__new__Nonecst|dk r||_nt|tr0t|j|_|dk rVtj|jkrNt d||_ n&tj|jkrr|jtj|_ n | |_ |dk r||_ n t|_ t|tjrt|ts||tj|jkrtj|jtj<dS)a Args: x: initial array for the MetaTensor. Can be a list, tuple, NumPy ndarray, scalar, and other types. affine: optional 4x4 array. meta: dictionary of metadata. applied_operations: list of previously applied operations on the MetaTensor, the list is typically maintained by `monai.transforms.TraceableTransform`. See also: :py:class:`monai.transforms.TraceableTransform` _args: additional args (currently not in use in this constructor). _kwargs: additional kwargs (currently not in use in this constructor). Note: If a `meta` dictionary is given, use it. Else, if `meta` exists in the input tensor `x`, use it. Else, use the default value. Similar for the affine, except this could come from four places, priority: `affine`, `meta["affine"]`, `x.affine`, `get_default_affine`. NzRSetting affine, but the applied meta contains an affine. This will be overwritten.)super__init__r2rrr__dict__rAFFINEwarningswarnr1get_default_affiner3get_default_applied_operationsrTensorrcopy_meta_fromSPACErRAS)selfr$r1r2r3_argsr< __class__rr r@vs&          zMetaTensor.__init__r)retsr4c Csg}d}tddt||D}t|D]l\}}t|tsBnNtsR|}n>t||} ||_ |j | | d|rt ||||||}| |q.t|t rt |S|S)a Update the metadata from the output of `MetaTensor.__torch_function__`. The output of `torch.Tensor.__torch_function__` could be a single object or a sequence of them. Hence, in `MetaTensor.__torch_function__` we convert them to a list of not already, and then we loop across each element, processing metadata as necessary. For each element, if not of type `MetaTensor`, then nothing to do. Args: rets: the output from `torch.Tensor.__torch_function__`, which has been converted to a list in `MetaTensor.__torch_function__` if it wasn't already a `Sequence`. func: the torch function that was applied. Examples might be `torch.squeeze` or `torch.Tensor.__add__`. We need this since the metadata need to be treated differently if a batch of data is considered. For example, slicing (`torch.Tensor.__getitem__`) the ith element of the 0th dimension of a batch of data should return a ith tensor with the ith metadata. args: positional arguments that were passed to `func`. kwargs: keyword arguments that were passed to `func`. Returns: A sequence with the same number of elements as `rets`. For each element, if the input type was not `MetaTensor`, then no modifications will have been made. If global parameters have been set to false (e.g., `not get_track_meta()`), then any `MetaTensor` will be converted to `torch.Tensor`. Else, metadata will be propagated as necessary (see :py:func:`MetaTensor._copy_meta`). Ncss|]}t|dr|jVqdS)is_batchN)rrPr"rrr r%s z)MetaTensor.update_meta..) copy_attr)r.rflatten_meta_objsvalues enumeraterrrr7rPrH_handle_batchedappendtuple) rOrr:r;outmetasrPidxr/Z meta_argsrrr update_metas   zMetaTensor.update_metac Cs|tjjkr|dks2t|dks2t|ddkr6|St|dtrP|ddn|d}|tdddtdfkszt|tjr~|St|ddd}||} t| t r| rz t | } Wqt t t tfk r} zt d| W5d} ~ XYqXnt| trd| _t| dr| j|_n|tjjkrt|dkr8|d} nd |krL|d } nd} | dkr|dkrtt|ddd}t||dr||j|_d|_|S) z/utility function to handle batched MetaTensors.rNF)detachzInconsistent batched metadata dicts when slicing a batch of MetaTensors, please consider converting it into a torch Tensor using `x.as_tensor()` or a numpy array using `x.array`.rAdim)rrG __getitem__lenrrsliceEllipsisr listr TypeError ValueError RuntimeError IndexErrorrrPrrAcopyunbind) r9r/rZrYrr:r; batch_idxZ dec_batchZret_metaer_rrr rUsD$""        zMetaTensor._handle_batchedrr)r4c s|dkr i}t||||}t|r*|St|dk rt|t|rt||||}t|jD]$}||j ||_ ||j ||_ q^|St|t s|g}d}nd}t||||}|r|dS|S)zWraps all torch functions.NTFr) r?__torch_function__r0r!rrr[rangen_fieldsr2r3r) r9rtypesr:r;r/Z out_itemsrZunpackrMrr rms" zMetaTensor.__torch_function__cCs,t|ttjttfr(t|tjdddS|S)NF) output_type wrap_sequencer) rrrrGrWrdrr,r-)r$rrr _convert2szMetaTensor._convertcCs\z|jdstWSWntk r.tYSXtttj|}dd|D}|||S)zQfor numpy Interoperability, so that we can compute ``np.sum(MetaTensor([1.0]))``.numpycSsi|]\}}|t|qSrrrtr#kvrrr @sz1MetaTensor.__array_function__..) __module__ startswithNotImplementedAttributeErrorrdmaprrtitems)rKrrpr:r;rLr<rrr __array_function__8s   zMetaTensor.__array_function__cOszt|jdstWSWntk r2tYSX|dkr@tSttj|}dd|D}d|krjtSzt ||||WStk rtYSXdS)z For numpy interoperability, so that we can compute ``MetaTensor([1.0]) >= np.asarray([1.0])``. This is for pytorch > 1.8. ru__call__cSsi|]\}}|t|qSrrvrwrrr rzPsz.MetaTensor.__array_ufunc__..rXN) rr{r|r}r~rrrtrr)rKufuncmethodinputsr;Z_inputsr<rrr __array_ufunc__Cs   zMetaTensor.__array_ufunc__z torch.TensorcCstjdtd|dS)Ncpur5)reyer+r*rrr rEXszMetaTensor.get_default_affinecCs |tjS)z Return the `MetaTensor` as a `torch.Tensor`. It is OS dependent as to whether this will be a deep copy or not. )r8rrGrKrrr r7\szMetaTensor.as_tensorcOst||||dddS)a Returns a new array in `output_type`, the array shares the same underlying storage when the output is a numpy array. Changes to self tensor will be reflected in the ndarray and vice versa. Args: output_type: output type, see also: :py:func:`monai.utils.convert_data_type`. dtype: dtype of output data. Converted to correct library type (e.g., `np.float32` is converted to `torch.float32` if output type is `torch.Tensor`). If left blank, it remains unchanged. device: if the output is a `torch.Tensor`, select device (if `None`, unchanged). _args: currently unused parameters. _kwargs: currently unused parameters. T)rrr*r+rsr)r)rKrrr*r+rLr<rrr get_arraycszMetaTensor.get_arrayFbool non_blockingcOsBt|ddd}z|j||dWStk r<||_|YSXdS)a Copies the elements from src into self tensor and returns self. The src tensor must be broadcastable with the self tensor. It may be of a different data type or reside on a different device. See also: `https://pytorch.org/docs/stable/generated/torch.Tensor.copy_.html` Args: src: the source tensor to copy from. non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously with respect to the host. For other cases, this argument has no effect. _args: currently unused parameters. _kwargs: currently unused parameters. FT) track_metarsrN)rcopy_rgdata)rKsrcrrLr< convertedrrr set_arrayss zMetaTensor.set_arraycCs|S)a Returns a numpy array of ``self``. The array and ``self`` shares the same underlying storage if self is on cpu. Changes to ``self`` (it's a subclass of torch.Tensor) will be reflected in the ndarray and vice versa. If ``self`` is not on cpu, the call will move the array to cpu and then the storage is not shared. :getter: see also: :py:func:`MetaTensor.get_array()` :setter: see also: :py:func:`MetaTensor.set_array()` )rrrrr arrays zMetaTensor.arraycCs||dS)z+A default setter using ``self.set_array()``N)r)rKrrrr rsr'dict)keyr4cCsJ|tjtjfkr td|d||j||dt||jt||j iS)a Get the object as a dictionary for backwards compatibility. This method does not make a deep copy of the objects. Args: key: Base key to store main data. The key for the metadata will be determined using `PostFix`. output_type: `torch.Tensor` or `np.ndarray` for the main data. dtype: dtype of output data. Converted to correct library type (e.g., `np.float32` is converted to `torch.float32` if output type is `torch.Tensor`). If left blank, it remains unchanged. Return: A dictionary consisting of three keys, the main data (stored under `key`) and the metadata. z4output_type must be torch.Tensor or np.ndarray, got .)rrr*) rrGr,r-rfrrr2 transformsr3)rKrrrr*rrr as_dicts zMetaTensor.as_dictcOst|tr,|dd^}}|s"|n|d}n t|dd}t|dddhdd}|dkr\tj}n|d krltj}nd }|j |||d S) a Cast to ``dtype``, sharing data whenever possible. Args: dtype: dtypes such as np.float32, torch.float, "np.float32", float. device: the device if `dtype` is a torch data type. _args: additional args (currently unused). _kwargs: additional kwargs (currently unused). Returns: data array instance rr]rr{rrur,)default)rur,N)rrr*r+) rr'splitrr rrGr,r-r)rKr*r+rLr<mod_strout_typerrr astypes  zMetaTensor.astypecCs|jtj|S)zAGet the affine. Defaults to ``torch.eye(4, dtype=torch.float64)``)r2getrrBrErrrr r1szMetaTensor.affiner)dr4cCs$tj|tdtjd|jtj<dS)zSet the affine.rr5N)rr7r+float64r2rrB)rKrrrr r1scCs |jrdd|jDSt|jS)zGet the spacingcSsg|] }t|qSr)r )r#arrr sz%MetaTensor.pixdim..)rPr1r rrrr pixdimszMetaTensor.pixdimcCsHd}|jr|jdtjd}|dkrDtt|jddddS|S)z Get the currently expected spatial shape as if all the pending operations are executed. For tensors that have more than 3 spatial dimensions, only the shapes of the top 3 dimensions will be returned. NT)rsr])pending_operationsrrSHAPErWrshapetolist)rKresrrr peek_pending_shapeszMetaTensor.peek_pending_shapecCs|j}t|d}|dkr,td|d|jD]T}t|tjt j d}|dkrVq2t ||d}t j j||}t jjj||}q2|S)Nr])r\z)Only 2d and 3d affine are supported, got zd input.rr)r1rarCrDrrrrrBrrrmonairutils to_affine_ndrlazyZcombine_transforms)rKrrpZ next_matrixrrr peek_pending_affines  zMetaTensor.peek_pending_affinecCsB|jr|jdtjdn|j}|dkr,dSttdt|dS)Nrr])rrrrBr1r&maxra)rKrrrr peek_pending_ranks zMetaTensor.peek_pending_rankcCst||j||||dS)z must be defined for deepcopy to work See: - https://pytorch.org/docs/stable/generated/torch.Tensor.new_empty.html#torch-tensor-new-empty )sizer*r+ requires_grad)rr7 new_empty)rKrr*r+rrrr rszMetaTensor.new_emptycKs$t|jf|}t|j|_|S)z Returns a copy of the MetaTensor instance. Args: kwargs: additional keyword arguments to `torch.clone`. See also: https://pytorch.org/docs/stable/generated/torch.clone.html )rr7clonerrA)rKr;Znew_instrrr rs  zMetaTensor.clonerz str | None)imr2 simple_keyspatternsepcCst|to|dk d}t|ts$|S|dkr0i}|rZtj|krRt|tj|tj<t||dk rxtjj ||dd|}|dkri}||_ tj|kr|tj|_ n t |_ |S)aX Convert the image to MetaTensor (when meta is not None). If `affine` is in the `meta` dictionary, convert that to `torch.Tensor`, too. Remove any superfluous metadata. Args: im: Input image (`np.ndarray` or `torch.Tensor`) meta: Metadata dictionary. When it's None, the metadata is not tracked, this method returns a torch.Tensor. simple_keys: whether to keep only a simple subset of metadata keys. pattern: combined with `sep`, a regular expression used to match and prune keys in the metadata (nested dictionary), default to None, no key deletion. sep: combined with `pattern`, used to match and delete keys in the metadata (nested dictionary). default is ".", see also :py:class:`monai.transforms.DeleteItemsd`. e.g. ``pattern=".*_code$", sep=" "`` removes any meta keys that ends with ``"_code"``. Returns: By default, a `MetaTensor` is returned. However, if `get_track_meta()` is `False` or meta=None, a `torch.Tensor` is returned. N)rT)keysrZuse_re) rrrrrrBr rrZ DeleteItemsdr2r1rE)rr2rrrimgrrr ensure_torch_and_prune_metas$    z&MetaTensor.ensure_torch_and_prune_metacCsd|S)z Prints a representation of the tensor. Prepends "meta" to ``torch.Tensor.__repr__``. Use ``print_verbose`` for associated metadata. r2)r7__repr__rrrr rGszMetaTensor.__repr__cCsdt|S)z Prints a representation of the tensor. Prepends "meta" to ``torch.Tensor.__str__``. Use ``print_verbose`` for associated metadata. r2)r'r7rrrr __str__OszMetaTensor.__str__cCs||S)zO returns the output of pytorch tensor's ``__format__`` method. )r7 __format__)rK format_specrrr rWszMetaTensor.__format__cCs$t||jdk r t|jdS)zVerbose print with meta data.N)printr2rrrrr print_verbose]s zMetaTensor.print_verbose)NNN)NNN)rN)F)N)NNF)FNr))rr{ __qualname____doc__ staticmethodr=r@r[ classmethodrUrmrtrrrrrEr7r,r-rrpropertyrsetterrGrrr1rrrrrrrrrrr __classcell__rrrMr r3sf576 0        2)) __future__r functoolsrCrirtypingrrrur,rrmonai.config.type_definitionsrmonai.data.meta_objrrmonai.data.utilsr r r r monai.utilsr monai.utils.enumsrrrrZmonai.utils.type_conversionrrrr__all__ lru_cacher!r0rGrrrrr  s$