U h@sddlmZmZmZmZmZmZmZddlZddlm Z ddl m Z ddl mZGdddeZGdd d eZGd d d eZGd d d eZdS))AnyCallableDictListOptionalSequenceUnionN)nn) transforms) TransformcsJeZdZdZeeddfdd ZeedddZe d d d Z Z S) ComposeaComposes several transforms together. This transform does not support torchscript. Please, see the note below. Args: transforms (list of ``Transform`` objects): list of transforms to compose. Example: >>> transforms.Compose([ >>> transforms.CenterCrop(10), >>> transforms.PILToTensor(), >>> transforms.ConvertImageDtype(torch.float), >>> ]) .. note:: In order to script the transformations, please use ``torch.nn.Sequential`` as below. >>> transforms = torch.nn.Sequential( >>> transforms.CenterCrop(10), >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), >>> ) >>> scripted_transforms = torch.jit.script(transforms) Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require `lambda` functions or ``PIL.Image``. Nr returncs4tt|tstdn |s*td||_dS)N5Argument transforms should be a sequence of callableszPass at least one transform)super__init__ isinstancer TypeError ValueErrorr selfr  __class__V/var/www/html/venv/lib/python3.8/site-packages/torchvision/transforms/v2/_container.pyr(s    zCompose.__init__inputsrcGs2t|dk}|jD]}||}|r&|n|f}q|SN)lenr rrneeds_unpacking transformoutputsrrrforward0s   zCompose.forwardrcCs*g}|jD]}|d|q d|SNz  r appendjoinr format_stringtrrr extra_repr7s zCompose.extra_repr) __name__ __module__ __qualname____doc__rrrrr$strr. __classcell__rrrrr sr csteZdZdZejZdeee e j fe ddfdd Z eeefddd Zeed d d Zedd dZZS) RandomApplyaApply randomly a list of transformations with a given probability. .. note:: In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of transforms as shown below: >>> transforms = transforms.RandomApply(torch.nn.ModuleList([ >>> transforms.ColorJitter(), >>> ]), p=0.3) >>> scripted_transforms = torch.jit.script(transforms) Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require `lambda` functions or ``PIL.Image``. Args: transforms (sequence or torch.nn.Module): list of transformations p (float): probability of applying the list of transforms ?Nr prcsPtt|ttjfs"td||_d|krs &  r5csHeZdZdZd eeeeeddfdd Z e e dddZ Z S) RandomChoiceaApply single transformation randomly picked from a list. This transform does not support torchscript. Args: transforms (sequence or torch.nn.Module): list of transformations p (list of floats or None, optional): probability of each transform being picked. If ``p`` doesn't sum to 1, it is automatically normalized. If ``None`` (default), all transforms have the same probability. Nr7cst|tstd|dkr*dgt|}n,t|t|krVtdt|dt|t||_t|fdd|D|_ dS)Nrrz4Length of p doesn't match the number of transforms: z != csg|] }|qSrr).0Zprobtotalrr sz)RandomChoice.__init__..) rrrrrrrr sumr8r:rrBrrs  zRandomChoice.__init__rcGs*ttt|jd}|j|}||Sr)intr=Z multinomialZtensorr8r )rridxr"rrrr$s zRandomChoice.forward)N) r/r0r1r2rrrrr?rrr$r4rrrrr@ts r@cs<eZdZdZeeddfdd ZeedddZZ S) RandomOrderzApply a list of transformations in a random order. This transform does not support torchscript. Args: transforms (sequence or torch.nn.Module): list of transformations Nr cs&t|tstdt||_dS)Nr)rrrrrr rrrrrs  zRandomOrder.__init__rcGsFt|dk}tt|jD]$}|j|}||}|r:|n|f}q|Sr)rr=Zrandpermr )rrr!rGr"r#rrrr$s   zRandomOrder.forward) r/r0r1r2rrrrr$r4rrrrrHsrH)typingrrrrrrrr=r Z torchvisionr r>Ztorchvision.transforms.v2r r r5r@rHrrrrs$   46%