o  i?O@sddlmZddlZddlZddlmZmZddlmZddl m Z ddl Z ddl m Z ddlmZmZddlmZdd lmZdd lmZdd lmZdd lmZmZmZmZmZmZm Z dd l!m"Z"ddgZ#GdddeZ$Gddde$eZ%dS)) annotationsN)HashableMapping)contextmanager)Any) transforms)MetaObjget_track_meta) MetaTensor) to_affine_nd)InvertibleTrait) Transform)LazyAttrMetaKeys TraceKeysTraceStatusKeysconvert_to_dst_typeconvert_to_numpyconvert_to_tensor) MONAIEnvVarsTraceableTransformInvertibleTransformc@seZdZdZddZddZed.dd Zejd/d d Ze d0d1ddZ e ddZ d2ddZ ddZ e d3d4ddZd5d"d#Zd6d7d'd(Zd8d9d)d*Zed:d,d-Zd S);ra. Maintains a stack of applied transforms to data. Data can be one of two types: 1. A `MetaTensor` (this is the preferred data type). 2. A dictionary of data containing arrays/tensors and auxiliary metadata. In this case, a key must be supplied (this dictionary-based approach is deprecated). If `data` is of type `MetaTensor`, then the applied transform will be added to ``data.applied_operations``. If `data` is a dictionary, then one of two things can happen: 1. If data[key] is a `MetaTensor`, the applied transform will be added to ``data[key].applied_operations``. 2. Else, the applied transform will be appended to an adjacent list using `trace_key`. If, for example, the key is `image`, then the transform will be appended to `image_transforms` (this dictionary-based approach is deprecated). Hopefully it is clear that there are three total possibilities: 1. data is `MetaTensor` 2. data is dictionary, data[key] is `MetaTensor` 3. data is dictionary, data[key] is not `MetaTensor` (this is a deprecated approach). The ``__call__`` method of this transform class must be implemented so that the transformation information is stored during the data transformation. The information in the stack of applied transforms must be compatible with the default collate, by only storing strings, numbers and arrays. `tracing` could be enabled by assigning to `self.tracing` or setting `MONAI_TRACE_TRANSFORM` when initializing the class. cCs8t|ds t|_t|jdstdk|j_dSdS)zRCreate a `_tracing` instance member to store the thread-local tracing state value._tracingvalue0N)hasattr threadinglocalrrtrace_transformrselfr!Z/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/transforms/inverse.py_init_trace_threadlocalJs   z*TraceableTransform._init_trace_threadlocalcsNttdi}fddtdgD}|ddt|dkr#|S||fS)zbWhen pickling, remove the `_tracing` member from the output, if present, since it's not picklable.__dict__csi|]}|t|qSr!)getattr).0krr!r" Xz3TraceableTransform.__getstate__.. __slots__rNr)dictr%poplen)r _dictZ_slotsr!rr" __getstate__Us zTraceableTransform.__getstate__returnboolcCs|t|jjS)z~ Returns the tracing state, which is thread-local and initialised to `MONAIEnvVars.trace_transform() != "0"`. )r#r1rrrr!r!r"tracing\s zTraceableTransform.tracingvalcCs|||j_dS)z-Sets the thread-local tracing state to `val`.N)r#rr)r r3r!r!r"r2ds NkeyrcCs|durtjS|tjS)z1The key to store the stack of applied transforms.N)r KEY_SUFFIX)r4r!r!r" trace_keyjszTraceableTransform.trace_keycCstjtjtjtjfS)z9The keys to store necessary info of an applied transform.)r CLASS_NAMEIDTRACING DO_TRANSFORMr!r!r!r"transform_info_keysqsz&TraceableTransform.transform_info_keysr+cCs8|jjt||jt|dr|jndf}tt||S)zg Return a dictionary with the relevant information pertaining to an applied transform. _do_transformT) __class____name__idr2rr<r+zipr;)r valsr!r!r"get_transform_infovs z%TraceableTransform.get_transform_infoc Os4|dd}|}|tjd}|pi}|dd}|rgtrgt|trg|sC|r0|j|ddni}|j ||tj |d} | | S|rT|j }| } ||n|i}} |j ||d| d} | | S||d<d|vr~t|dtr~|d|n||d<tj|g|Ri|} t|tr| | S|S) a Push to a stack of applied transforms of ``data``. Args: data: dictionary of data or `MetaTensor`. args: additional positional arguments to track_transform_meta. kwargs: additional keyword arguments to track_transform_meta, set ``replace=True`` (default False) to rewrite the last transform infor in applied_operation/pending_operation based on ``self.get_transform_info()``. lazyFTreplace)check) orig_size extra_info)transform_inforCrGrH)getrBrr:r,r isinstancer pop_transformpush_transform ORIG_SIZEcopy_meta_frompending_operationscopyupdater+rtrack_transform_meta) r dataargskwargsZ lazy_evalrHZ do_transformrDZxformmeta_objextrar!r!r"rLs.       z!TraceableTransform.push_transformFrG dict | NonerF tuple | Nonec  Cs|dur||n|} t} t| tr| j| | jd|r&ts&td|s~|dur~t| tr~| } t | |t j dd} z| t t| d|t j d}Wntyn} z| jdkri| jrbd} nd } t| | d} ~ wwt|t d t j d | jtj<tr|r|tjst|trt|tst|}t| tr| | n| ||<|S| S|}|dur||tj<nt| tr| |tj<nt| d r| j dd|tj<||tj!<|dur|"t#j$d|"t#jd||tj%<|rJ|durt#j$|vr|tjg|t#j$<n||t#j$<t&t'|t#j$d d(|t#j$<|dur0t#j|vr/t)|t#j<n||t#j<t|t#jt d d|t#j<| *|nV| j+rt|tr[|tj,dnd}d|ddd| j+D} |durw| d|7} | j+d}|tj-t}|t.j/t0}|1| ||t.j/<||tj-<| 2|t|trt|tst|}t| tr| | ||<|St34|}||vrg||<||1||S| S)am Update a stack of applied/pending transforms metadata of ``data``. Args: data: dictionary of data or `MetaTensor`. key: if data is a dictionary, data[key] will be modified. sp_size: the expected output spatial size when the transform is applied. it can be tensor or numpy, but will be converted to a list of integers. affine: the affine representation of the (spatial) transform in the image space. When the transform is applied, meta_tensor.affine will be updated to ``meta_tensor.affine @ affine``. extra_info: if desired, any extra information pertaining to the applied transform can be stored in this dictionary. These are often needed for computing the inverse transformation. orig_size: sometimes during the inverse it is useful to know what the size of the original image was, in which case it can be supplied here. transform_info: info from self.get_transform_info(). lazy: whether to push the transform to pending_operations or applied_operations. Returns: For backward compatibility, if ``data`` is a dictionary, it returns the dictionary with updated ``data[key]``. Otherwise, this function returns a MetaObj with updated transform metadata. N)keyszUmetadata is not tracked, please call 'set_track_meta(True)' if doing lazy evaluation.)dtyperzHTransform applied to batched tensor, should be applied to instances onlyz\Mismatch affine matrix, ensured that the batch dimension is not included in the calculation.cpu)devicer[shapeT) wrap_sequence)r_z Transform z; has been applied to a MetaTensor with pending operations: cSsg|]}|tjqSr!)rIrr7)r&xr!r!r" r)z;TraceableTransform.track_transform_meta..z for key )5rrJr rNr$rZr warningswarnpeek_pending_affinertorchfloat64r r- RuntimeErrorndimis_batchrr_metarAFFINErIrr9rr+rPrMpeek_pending_shaperr`LAZYr,rSHAPE EXTRA_INFOtuplertolistget_default_affinepush_pending_operationrOr7STATUSESrPENDING_DURING_APPLYlistappendpush_applied_operationrr6)clsrSr4Zsp_sizeaffinerGrFrHrCZdata_tZout_objZ orig_affineemsginfoZtransform_nameZpendstatusesmessagesZx_kr!r!r"rRs#                                z'TraceableTransform.track_transform_meta transformrNonec Cs|tjd}|t|krdS|tjkrdS|tjd}|tjid}|r.t|t j dvr=||j j kr=dStd|j j d|d|dt|d ) z&Check transforms are of same instance.rbNrg)spawnNzError z8 getting the most recently applied invertible transform  z != .)rIrr8r?NONEr7rsrfrgrimultiprocessingget_start_methodr=r>rk)r rZxform_idZ xform_name warning_msgr!r!r"check_transforms_match+s(    z)TraceableTransform.check_transforms_matchTrEr,cCs|jstdt|tr|j}n,t|tr2||vr&t||tr&||j}n|||t}n t dt |d|sNt dt |d|d|d|rW| |d|r^| dS|dS) aL Get most recent matching transform for the current class from the sequence of applied operations. Args: data: dictionary of data or `MetaTensor`. key: if data is a dictionary, data[key] will be modified. check: if true, check that `self` is the same type as the most recently-applied transform. pop: if true, remove the transform as it is returned. Returns: Dictionary of most recently applied transform Raises: - RuntimeError: data is neither `MetaTensor` nor dictionary zCTransform Tracing must be enabled to get the most recent transform.z8`data` should be either `MetaTensor` or dictionary, got rz Item of type z (key: z, pop: z ) has empty 'applied_operations're) r2rkrJr applied_operationsrrIr6get_default_applied_operations ValueErrortyperr,)r rSr4rEr,all_transformsr!r!r"get_most_recent_transform?s    z,TraceableTransform.get_most_recent_transformcCs|j|||ddS)a Return and pop the most recent transform. Args: data: dictionary of data or `MetaTensor` key: if data is a dictionary, data[key] will be modified check: if true, check that `self` is the same type as the most recently-applied transform. Returns: Dictionary of most recently applied transform Raises: - RuntimeError: data is neither `MetaTensor` nor dictionary T)r,)r)r rSr4rEr!r!r"rKcsz TraceableTransform.pop_transformto_traceccs|j}||_dV||_dS)zITemporarily set the tracing status of a transform with a context manager.N)r2)r rprevr!r!r"rts  z"TraceableTransform.trace_transform)r0r1)r3r1)N)r4r)r0r+)NNNNNNF)r4rrGrXrFrY)rrr0r)NTF)r4rrEr1r,r1)NT)r4rrEr1)rr1)r> __module__ __qualname____doc__r#r/propertyr2setter staticmethodr6r;rBrL classmethodrRrrrKrrr!r!r!r"r*s:      %   $c@s"eZdZdZddZd ddZd S) ra:Classes for invertible transforms. This class exists so that an ``invert`` method can be implemented. This allows, for example, images to be cropped, rotated, padded, etc., during training and inference, and after be returned to their original size before saving to file for comparison in an external viewer. When the ``inverse`` method is called: - the inverse is called on each key individually, which allows for different parameters being passed to each label (e.g., different interpolation for image and label). - the inverse transforms are applied in a last-in-first-out order. As the inverse is applied, its entry is removed from the list detailing the applied transformations. That is to say that during the forward pass, the list of applied transforms grows, and then during the inverse it shrinks back down to an empty list. We currently check that the ``id()`` of the transform is the same in the forward and inverse directions. This is a useful check to ensure that the inverses are being processed in the correct order. Note to developers: When converting a transform to an invertible transform, you need to: #. Inherit from this class. #. In ``__call__``, add a call to ``push_transform``. #. Any extra information that might be needed for the inverse can be included with the dictionary ``extra_info``. This dictionary should have the same keys regardless of whether ``do_transform`` was `True` or `False` and can only contain objects that are accepted in pytorch data loader's collate function (e.g., `None` is not allowed). #. Implement an ``inverse`` method. Make sure that after performing the inverse, ``pop_transform`` is called. cCsdt|tr t|tjs |St|}||D]}tj|}||vs&||s'qtj||dd}q|S)z This function is to be called before every `self.inverse(data)`, update each MetaTensor `data[key]` using `data[key_transforms]` and `data[key_meta_dict]`, for MetaTensor backward compatibility 0.9.0. F)t)rJr+r MapTransformZ key_iteratorrr6sync_meta_info)r rSdr'Z transform_keyr!r!r"inverse_updates z"InvertibleTransform.inverse_updaterSrr0cCstd|jjd)z Inverse of ``__call__``. Raises: NotImplementedError: When the subclass does not override this method. z Subclass z must implement this method.)NotImplementedErrorr=r>)r rSr!r!r"inverseszInvertibleTransform.inverseN)rSrr0r)r>rrrrrr!r!r!r"r}s$)& __future__rrrfcollections.abcrr contextlibrtypingrrimonairmonai.data.meta_objrr monai.data.meta_tensorr monai.data.utilsr monai.transforms.traitsr monai.transforms.transformr monai.utilsrrrrrrrmonai.utils.miscr__all__rrr!r!r!r"s(        $ U