o # il@s$ddlmZddlZddlZddlmZddlmZddlm Z ddl Z ddl Z ddl Z ddlmZddlmZmZddlmZmZmZmZdd lmZdd lmZmZmZmZdd lm Z m!Z!m"Z"m#Z#d gZ$e%dd dZ&ddZ'Gdd d ee j(Z)e*e j+dre j+,ee)eegdSdS)) annotationsN)Sequence)deepcopy)Any) 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 MetaTensorcCsHttdr"t|dr"ttj|jr"tttj|jtr"ttj|jSdS)N return_types__name__)hasattrtorchrr isinstancegetattrtype)funcrX/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/data/meta_tensor.py_get_named_tuple_like_type"s r!c CsHt|ttttjtjtjtj fp#t|t p"t|t o"t dd|D S)Ncss|]}t|tVqdSN)rr.0xrrr 0sz*_not_requiring_metadata..) rintstrbytesrSizedtypedevicenpndarrayrrany)retrrr _not_requiring_metadata.s $r1cseZdZdZe   dVdWd d Z   dVdXfd d ZedYddZeddZ edZd[fdd Z eddZ ddZ ddZ eejfd\d d!Zd\d"d#Zejddfd$d%Zd]d^d)d*Zed+d,Zejd_d-d,Zejdfd`d1d2Zdad3d4Zed\d5d6Zejdbd9d6Zed:d;Zdd?Zd@dAZ dcdBdCZ!dDdEZ"e FdddedLdMZ#dNdOZ$dPdQZ%dRdSZ&d_dTdUZ'Z(S)fra 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. Naffinetorch.Tensor | Nonemeta dict | Noneapplied_operations list | NonereturncOs@|r|dd|dddni}tj|g|Ri||S)Nr,r+r,r+)popr as_tensor as_subclass)clsr%r2r4r6argskwargs_kwargsrrr __new__js" zMetaTensor.__new__Nonecst|dur ||_n t|trt|j|_|dur+tj|jvr't d||_ ntj|jvr9|jtj|_ n| |_ |durF||_ nt|_ t|tjr[t|ts[||tj|jvrjtj|jtj<dSdS)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__r4rrr__dict__rAFFINEwarningswarnr2get_default_affiner6get_default_applied_operationsrTensorrcopy_meta_fromSPACErRAS)selfr%r2r4r6_argsr@ __class__rr rDws(          zMetaTensor.__init__retsrc Csg}d}tddt||D}t|D]6\}}t|ts!n'ts)|}nt||} ||_ |j | | d|rHt ||||||}| |qt|t rWt |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)rrTr#rrr r&sz)MetaTensor.update_meta..) copy_attr)r/rflatten_meta_objsvalues enumeraterrrr;rTrL_handle_batchedappendtuple) rSrr>r?outmetasrTidxr0Z meta_argsrrr update_metas   zMetaTensor.update_metac Cs|tjjkr|dkst|dkst|ddkr|St|dtr'|ddn|d}|tdddtdfvs|St|ddd}||} t| t rm| rmzt | } Wnt t t tfyl} zt d| d} ~ wwt| trud| _t| dr| j|_|S|tjjkrt|dkr|d} n d |vr|d } nd} | dkr|durt|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`.rEdim)rrK __getitem__lenrrsliceEllipsisr listr TypeError ValueError RuntimeError IndexErrorrrTrrEcopyunbind) r=r0r^r]rr>r? batch_idxZ dec_batchZret_metaercrrr rYsJ $""       zMetaTensor._handle_batchedrrc s|duri}t||||}t|r|St|durDt|t|rDt||||}t|jD]}||j ||_ ||j ||_ q/|St|t sO|g}d}nd}t||||}|r_|dS|S)zWraps all torch functions.NTFr) rC__torch_function__r1r!rrr_rangen_fieldsr4r6r) r=rtypesr>r?r0Z out_itemsr^unpackrQrr rqs" zMetaTensor.__torch_function__cCs,t|ttjttfrt|tjdddS|S)NF) output_type wrap_sequencer) rrrrKr[rhrr-r.)r%rrr _convert3szMetaTensor._convertcCs^z |jds tWSWn tytYSwtttj|}dd|D}||i|S)zQfor numpy Interoperability, so that we can compute ``np.sum(MetaTensor([1.0]))``.numpycSi|] \}}|t|qSrrrxr$kvrrr Az1MetaTensor.__array_function__..) __module__ startswithNotImplementedAttributeErrorrhmaprrxitems)rOrrtr>r?rPr@rrr __array_function__9s  zMetaTensor.__array_function__cOsz t|jds tWSWn tytYSw|dkrtSttj|}dd|D}d|vr4tSz t |||i|WStyJtYSw)z For numpy interoperability, so that we can compute ``MetaTensor([1.0]) >= np.asarray([1.0])``. This is for pytorch > 1.8. ry__call__cSrzrr{r|rrr rQrz.MetaTensor.__array_ufunc__..r\) rrrrrrrrxrr)rOufuncmethodinputsr?Z_inputsr@rrr __array_ufunc__Ds$   zMetaTensor.__array_ufunc__ torch.TensorcCstjdtd|dS)Ncpur9)reyer,r+rrr rIYszMetaTensor.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. )r<rrKrOrrr r;]s zMetaTensor.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)rvr+r,rwr)r)rOrvr+r,rPr@rrr get_arraydszMetaTensor.get_arrayF non_blockingboolcOs<t|ddd}z|j||dWSty||_|YSw)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_metarw)r)rcopy_rkdata)rOsrcrrPr@ convertedrrr set_arrayts 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)rOrrrr rskeyr(dictcCsJ|tjtjfvrtd|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 .)rvr+) rrKr-r.rjrrr4 transformsr6)rOrrvr+rrr as_dicts   zMetaTensor.as_dictcOs~t|tr|dd^}}|s|n|d}nt|dd}t|hddd}|dkr-tj}n |d vr5tj}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 rrarrr>ryrr-ry)default)ryr-N)rvr+r,) rr(splitrr rrKr-r.r)rOr+r,rPr@mod_strout_typerrr astypes  zMetaTensor.astypecCs|jtj|S)zAGet the affine. Defaults to ``torch.eye(4, dtype=torch.float64)``)r4getrrFrIrrrr r2szMetaTensor.affinedrcCs$tj|tdtjd|jtj<dS)zSet the affine.rr9N)rr;r,float64r4rrF)rOrrrr r2s$cCs |jr dd|jDSt|jS)zGet the spacingcSsg|]}t|qSr)r )r$arrr sz%MetaTensor.pixdim..)rTr2r rrrr pixdims zMetaTensor.pixdimcCsHd}|jr|jdtjd}|dur"tt|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)rwra)pending_operationsrrSHAPEr[rshapetolist)rOresrrr peek_pending_shapes*zMetaTensor.peek_pending_shapecCs|j}t|d}|dvrtd|d|jD]*}t|tjt j d}|dur+qt ||d}t j j||}t jjj||}q|S)Nra)r`z)Only 2d and 3d affine are supported, got zd input.rr)r2rerGrHrrrrrFrrrmonairutils to_affine_ndrlazyZcombine_transforms)rOrrpZ next_matrixrrr peek_pending_affines  zMetaTensor.peek_pending_affinecCsB|jr |jdtjdn|j}|durdSttdt|dS)Nrra)rrrrFr2r'maxre)rOrrrr 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)rr; new_empty)rOrr+r,rrrr rszMetaTensor.new_emptycKs(t|jdi|}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 Nr)rr;clonerrE)rOr?Znew_instrrr rs zMetaTensor.clonerim simple_keyspattern str | NonesepcCst|to|dud}t|ts|S|duri}|r-tj|vr)t|tj|tj<t||durs6         5