o
    i                     @  s   d dl mZ d dlmZ d dlZd dlmZ d dlm	Z	 d dl
mZ d dlmZ d dlmZ d	d
lmZmZ ddlmZmZmZ g dZG dd deeZG dd deeZG dd deeZe ZZe ZZe ZZ dS )    )annotations)HashableN)KeysCollection)NdarrayOrTensor)get_track_meta)convert_to_tensor)ensure_tuple   )MapTransformRandomizableTransform   )CutMixCutOutMixUp)	MixUpdMixUpD	MixUpDictCutMixdCutMixD
CutMixDictCutOutdCutOutD
CutOutDictc                      s:   e Zd ZdZ	ddddZdd fddZdd Z  ZS )r   z
    Dictionary-based version :py:class:`monai.transforms.MixUp`.

    Notice that the mixup transformation will be the same for all entries
    for consistency, i.e. images and labels must be applied the same augmenation.
          ?Fkeysr   
batch_sizeintalphafloatallow_missing_keysboolreturnNonec                 C  s   t | || t||| _d S N)r
   __init__r   mixup)selfr   r   r   r    r'   l/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/transforms/regularization/dictionary.pyr$   &   s   zMixUpd.__init__Nseed
int | Nonestatenp.random.RandomState | Nonec                       t  || | j|| | S r#   )superset_random_stater%   r&   r)   r+   	__class__r'   r(   r/   ,      zMixUpd.set_random_statec                 C  s>   t |}| jd  | |D ]}| j|| dd||< q|S )NF	randomize)dictr%   r5   key_iterator)r&   datadkr'   r'   r(   __call__1   s
   zMixUpd.__call__)r   F)
r   r   r   r   r   r   r   r    r!   r"   NN)r)   r*   r+   r,   r!   r   __name__
__module____qualname____doc__r$   r/   r;   __classcell__r'   r'   r1   r(   r      s    r   c                      sB   e Zd ZdZ			dd fddZdd fddZdd Z  ZS )r   a  
    Dictionary-based version :py:class:`monai.transforms.CutMix`.

    Notice that the mixture weights will be the same for all entries
    for consistency, i.e. images and labels must be aggregated with the same weights,
    but the random crops are not.
    Nr   Fr   r   r   r   
label_keysKeysCollection | Noner   r   r   r    r!   r"   c                   s:   t  || t||| _|d urt|| _d S g | _d S r#   )r.   r$   r   mixerr   rC   )r&   r   r   rC   r   r   r1   r'   r(   r$   C   s    zCutMixd.__init__r)   r*   r+   r,   c                   r-   r#   )r.   r/   rE   r0   r1   r'   r(   r/   O   r3   zCutMixd.set_random_statec                 C  s   t |}| |}|dkrt|t d}|S | j||  | || jD ]"\}}| j|| ||d dd}|d ||< ||v rG|d ||< q%|S )Nr'   
track_metaFr4   r   r   )	r6   	first_keyr   r   rE   r5   r7   rC   get)r&   r8   r9   rH   outkey	label_keyretr'   r'   r(   r;   T   s   
zCutMixd.__call__)Nr   F)r   r   r   r   rC   rD   r   r   r   r    r!   r"   r<   )r)   r*   r+   r,   r!   r   r=   r'   r'   r1   r(   r   :   s    r   c                      s<   e Zd ZdZdd fddZdd fddZdd Z  ZS )r   z
    Dictionary-based version :py:class:`monai.transforms.CutOut`.

    Notice that the cutout is different for every entry in the dictionary.
    Fr   r   r   r   r   r    r!   r"   c                   s   t  || t|| _d S r#   )r.   r$   r   cutout)r&   r   r   r   r1   r'   r(   r$   j   s   zCutOutd.__init__Nr)   r*   r+   r,   c                   r-   r#   )r.   r/   rN   r0   r1   r'   r(   r/   n   r3   zCutOutd.set_random_statec                 C  sf   t |}| |}|dkrt|t d}|S | j||  | |D ]}| j|| dd||< q#|S )Nr'   rF   Fr4   )r6   rH   r   r   rN   r5   r7   )r&   r8   r9   rH   rJ   r:   r'   r'   r(   r;   s   s   
zCutOutd.__call__)F)r   r   r   r   r   r    r!   r"   r<   )r)   r*   r+   r,   r!   r   r=   r'   r'   r1   r(   r   c   s
    r   )!
__future__r   collections.abcr   numpynpmonai.configr   monai.config.type_definitionsr   monai.data.meta_objr   monai.utilsr   monai.utils.miscr   	transformr
   r   arrayr   r   r   __all__r   r   r   r   r   r   r   r   r   r'   r'   r'   r(   <module>   s"   )