a h@n@sddlZddlmZddlmZddlZddlmZddlmZ m Z gdZ ee e e eee dd d ZGd d d eZGd ddejjZGdddejjZGdddejjZGdddejjZdS)N)Enum)Optional)Tensor) functionalInterpolationMode)AutoAugmentPolicy AutoAugment RandAugmentTrivialAugmentWideAugMiximgop_name magnitude interpolationfillc Cs|dkr>tj|dddgdtt|dg||ddgd}n|dkr|tj|dddgddtt|g||ddgd}n|dkrtj|dt|dgd|ddg|d}nP|d krtj|ddt|gd|ddg|d}n |d krtj||||d }n|d krt|d|}n|d kr2t|d|}n|dkrNt |d|}n|dkrjt |d|}n|dkrt |t|}nv|dkrt ||}n^|dkrt |}nH|dkrt|}n2|dkrt|}n|dkrntd|d|S)NShearXr?)angle translatescaleshearrrcenterShearY TranslateX)rrrrrr TranslateYRotaterr BrightnessColorContrast Sharpness PosterizeSolarize AutoContrastEqualizeInvertIdentityzThe provided operator  is not recognized.)FZaffinemathdegreesatanintrotateZadjust_brightnessZadjust_saturationZadjust_contrastZadjust_sharpnessZ posterizeZsolarizeZ autocontrastZequalizeinvert ValueErrorr r3P/var/www/auris/lib/python3.9/site-packages/torchvision/transforms/autoaugment.py _apply_op s                   r5c@seZdZdZdZdZdZdS)rzoAutoAugment policies learned on different datasets. Available policies are IMAGENET, CIFAR10 and SVHN. ZimagenetZcifar10ZsvhnN)__name__ __module__ __qualname____doc__IMAGENETCIFAR10SVHNr3r3r3r4r]src seZdZdZejejdfeeee e ddfdd Z ee e e e e eefe e e eeffdddZee eefee e eeffd d d Zeee eeefd d dZeedddZe dddZZS)r a?AutoAugment data augmentation method based on `"AutoAugment: Learning Augmentation Strategies from Data" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: policy (AutoAugmentPolicy): Desired policy enum defined by :class:`torchvision.transforms.autoaugment.AutoAugmentPolicy`. Default is ``AutoAugmentPolicy.IMAGENET``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. N)policyrrreturncs,t||_||_||_|||_dSN)super__init__r=rr _get_policiespolicies)selfr=rr __class__r3r4rAys  zAutoAugment.__init__)r=r>cCsJ|tjkrgdS|tjkr$gdS|tjkr6gdStd|ddS)N)))r$皙?)r333333? )r%rIr&rINr'皙?Nr'rIN))r$rI)r$rIr'rGN)r%皙?)rUrrPrH))r%rIrQ))r$rPrLr'rN))rrVrY)r%rIrH)rQ)r$rGrS)rXr!rGr))rrGrJrQ))r'rNrOr(rINrZ)r!rIrW)r"rrH)rX)r!r))r!rPrH)r%rPrR))r#rGrRr]))rrIrLrZ)r[rQrTrKr\r^rN))r(皙?N)r"rVrS))rffffff?r_)r333333?rJ))r#rPr)r#?rY))r?rHrrbrJ))r&reNr'rdN))rrVrR)r$rcrR))r!rGrY)r rIrR))r#rcrJ)r rbrJ)rQ)r'reN))r"rIrR)r#rIrL))r!rbrR)rrerH))r'rcN)r&rGN))rrGrY)r#rVrS))r rdrS)r!rVrH))r%rer_)r(rN)r'rVNrM)rhrQ))r!rdrJrQ)r&rPN)r%rVrH))r rarY)r!rbr))r%rGrLr&rdN))rrdrJrf)rj)r%rPrY)rOr`)rfrj))rrdrW)r(rVN)rrdrHr(rbN)rQ)r%rIrSr(rdNrQrQ)rrdrY)rkri)rl)r(rGN))rrdrL)r%rVrS)rorirp)rk)r%rcrY))rrPrHrm)rg)rrIrSrn))r"rcrYrrPrW)r(rPN)rrr_))rrbrS)r%rGrH)r]rq))rrcrR)rrdrY))rrarSr]))r%rbr_)rrIrR))rrPrWrr))rrbrJ)rrPrY))rrPrL)r&rbN))rrbr_r`zThe provided policy r*)rr:r;r<r2)rDr=r3r3r4rBs   zAutoAugment._get_policiesnum_bins image_sizer>cCstdd|dftdd|dftdd|d|dftdd|d|dftdd|dftdd|dftdd|dftdd|dftdd|dfd t||dd d ftd d|d ftdd ftdd ftdd fd S)NrrcTt?rr>@rdrHrWFo@)rrrrrr r!r"r#r$r%r&r'r()torchlinspacearangeroundr/tensorrDrtrur3r3r4_augmentation_spaces$   zAutoAugment._augmentation_space) transform_numr>cCs4tt|d}td}tdd}|||fS)zGet parameters for autoaugment transformation Returns: params required by the autoaugment transformation r)r_r_)r/ryrandintitemZrand)rZ policy_idprobssignsr3r3r4 get_paramss  zAutoAugment.get_paramsrr>cCs|j}t|\}}}t|trTt|ttfr>t|g|}n|durTdd|D}|t|j \}}}| d||f} t |j |D]n\} \} } } || | kr| | \}}| durt||  nd}|r|| dkr|d9}t || ||j|d}q|S) z img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: AutoAugmented image. NcSsg|] }t|qSr3float.0fr3r3r4 z'AutoAugment.forward.. rrr)rr+get_dimensions isinstancerr/rrlenrCr enumeraterr5r)rDrrchannelsheightwidthZ transform_idrrop_metairpZ magnitude_id magnitudessignedrr3r3r4forwards"   zAutoAugment.forwardr>cCs|jjd|jd|jdS)Nz(policy=, fill=))rFr6r=r)rDr3r3r4__repr__szAutoAugment.__repr__)r6r7r8r9rr:rNEARESTrlistrrAtuplestrr/rBdictrboolr staticmethodrrr __classcell__r3r3rEr4r hs$  * Z* r cseZdZdZdddejdfeeeeeee ddfdd Z ee eefe e e eeffd d d Zeed d dZe dddZZS)r a~RandAugment data augmentation method based on `"RandAugment: Practical automated data augmentation with a reduced search space" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: num_ops (int): Number of augmentation transformations to apply sequentially. magnitude (int): Magnitude for all the transformations. num_magnitude_bins (int): The number of different magnitude values. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. r_rJN)num_opsrnum_magnitude_binsrrr>cs,t||_||_||_||_||_dSr?)r@rArrrrr)rDrrrrrrEr3r4rA2s  zRandAugment.__init__rscCstddftdd|dftdd|dftdd|d|dftdd|d|dftdd|dftdd |dftdd |dftdd |dftdd |dfd t||dd dftd d|dftddftddfd S)NrFrcTrvrrrwrdrHrWrxr)rrrrrr r!r"r#r$r%r&r'ryr}rzr{r|r/r~r3r3r4rAs $  zRandAugment._augmentation_spacerc Cs|j}t|\}}}t|trTt|ttfr>t|g|}n|durTdd|D}||j||f}t |j D]}tt t |d}t||} || \} } | jdkrt| |jnd} | rt ddr| d9} t|| | |j|d }qp|S) img (PIL Image or Tensor): Image to be transformed. Returns: PIL Image or Tensor: Transformed image. NcSsg|] }t|qSr3rrr3r3r4rarz'RandAugment.forward..rrrr_rr)rr+rrrr/rrrrangerryrrrrkeysndimrr5r) rDrrrrrr_op_indexrrrrr3r3r4rTs"   zRandAugment.forwardrc Cs:|jjd|jd|jd|jd|jd|jd }|S)Nz (num_ops=z , magnitude=z, num_magnitude_bins=, interpolation=rr)rFr6rrrrrrDsr3r3r4ros  zRandAugment.__repr__)r6r7r8r9rrr/rrrrArrrrrrrrrr3r3rEr4r s" *r cs|eZdZdZdejdfeeeee ddfdd Z ee e e eeffddd Zeed d d Ze d ddZZS)r aDataset-independent data-augmentation with TrivialAugment Wide, as described in `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: num_magnitude_bins (int): The number of different magnitude values. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. rN)rrrr>cs t||_||_||_dSr?)r@rArrr)rDrrrrEr3r4rAs zTrivialAugmentWide.__init__)rtr>cCstddftdd|dftdd|dftdd|dftdd|dftdd|dftdd|dftdd|dftdd|dftdd|dfdt||dd dftd d|dftddftddfd S) NrFgGz?Tg@@g`@rHrrSrxrr)rDrtr3r3r4rs $  z&TrivialAugmentWide._augmentation_spacerc Cs|j}t|\}}}t|trTt|ttfr>t|g|}n|durTdd|D}||j}tt t |d }t ||}||\} } | jdkrt| t j t | dt jd nd} | rt ddr| d 9} t||| |j|d S) rNcSsg|] }t|qSr3rrr3r3r4rrz.TrivialAugmentWide.forward..rrdtyperr_rr)rr+rrrr/rrrryrrrrrrlongr5r) rDrrrrrrrrrrrr3r3r4rs$   $zTrivialAugmentWide.forwardrcCs*|jjd|jd|jd|jd}|S)Nz(num_magnitude_bins=rrr)rFr6rrrrr3r3r4rs zTrivialAugmentWide.__repr__)r6r7r8r9rrr/rrrrArrrrrrrrrr3r3rEr4r |s  r c seZdZdZdddddejdfeeeeeee e eddfdd Z ee eefe ee eeffd d d Zejjed ddZejjedddZeedddZeedddZed ddZZS)r aAugMix data augmentation method based on `"AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty" `_. If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, it is expected to be in mode "L" or "RGB". Args: severity (int): The severity of base augmentation operators. Default is ``3``. mixture_width (int): The number of augmentation chains. Default is ``3``. chain_depth (int): The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3]. Default is ``-1``. alpha (float): The hyperparameter for the probability distributions. Default is ``1.0``. all_ops (bool): Use all operations (including brightness, contrast, color and sharpness). Default is ``True``. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported. fill (sequence or number, optional): Pixel fill value for the area outside the transformed image. If given a number, the value is used for all bands respectively. rYrTN)severity mixture_width chain_depthalphaall_opsrrr>csntd|_d|kr&|jks@ntd|jd|d||_||_||_||_||_||_ ||_ dS)Nrrz!The severity must be between [1, z]. Got z instead.) r@rA_PARAMETER_MAXr2rrrrrrr)rDrrrrrrrrEr3r4rAs zAugMix.__init__rsc Cs tdd|dftdd|dftd|dd|dftd|dd|dftdd|dfdt||ddd ftd d|d ftdd ftdd fd }|jr|tdd |dftdd |dftdd |dftdd |dfd |S)NrrcTrg@rrwrWFrx) rrrrrr$r%r&r'rd)r r!r"r#)ryrzr{r|r/r}rupdate)rDrtrurr3r3r4rs&$   zAugMix._augmentation_spacercCs t|Sr?)r+Z pil_to_tensorrDrr3r3r4_pil_to_tensorszAugMix._pil_to_tensor)rcCs t|Sr?)r+Z to_pil_imagerr3r3r4_tensor_to_pilszAugMix._tensor_to_pil)paramsr>cCs t|Sr?)ry_sample_dirichlet)rDrr3r3r4rszAugMix._sample_dirichlet)orig_imgr>c Cs||j}t|\}}}t|trZ|}t|ttfrBt|g|}qd|durddd|D}n ||}||j ||f}t |j }| dgt d|jd|} | dgdg| jd} |tj|j|jg| jd| dd} |tj|jg|j| jd| dd| dddf | ddg} | dddf | | } t|jD]}| }|jdkrn|jnttjddd d }t|D]}ttt|d }t ||}||\}}|jdkrt|tj|jd tjd nd }|rtd d r|d9}t||||j |d}q| !| dd|f | |qT| |j"|j#d } t|tsx|$| S| S)rNcSsg|] }t|qSr3rrr3r3r4r/rz"AugMix.forward..rrWr)devicerr)lowhighsizerrr_rr)%rr+rrrr/rrrrrshapeviewmaxrrrryr}rrexpandrrrrrrrrrr5rZadd_torr)rDrrrrrrrZ orig_dimsbatchZ batch_dimsmZcombined_weightsZmixraugdepthrrrrrrr3r3r4r!sR    "$*  $$  zAugMix.forwardcCsJ|jjd|jd|jd|jd|jd|jd|jd|jd}|S) Nz (severity=z, mixture_width=z, chain_depth=z, alpha=z , all_ops=rrr) rFr6rrrrrrrrr3r3r4r[s"  zAugMix.__repr__)r6r7r8r9rZBILINEARr/rrrrrArrrrrryZjitZunusedrrrrrrr3r3rEr4r s4 *:r )r,enumrtypingrryrrr+r__all__rrrr5rnnModuler r r r r3r3r3r4s    P 8]V