khS?6dZddlZddlmZmZmZddlZddlmZddl m Z m Z m Z ddl mZmZmZddlmZmZmZmZmZmZmZmZmZddlmZmZmZdd lm Z d d e e d d fdee!ee!e!ffde"dee#dfdee#dfde$de$f dZ%d dddddddd dddde e ddddd d d d ddd fdee!ee!e!ffdeee#e#fdeee#e#fd ee"d!e#d"e#d#ee#ee#dffd$ee#d%e$d&e#d'e#d(ee"de"dee#dfdee#dfd)e#d*e"d+e!d,e!de$de$d-e$d.e$d/ee!ee!e!ffd0e!d1e$f4d2Z&d dddd e e d d d ddd f dee!ee!e!ffd3ee#d4ee"d5ee!de"dee#dfdee#dfde$de$d.e$d/ee!ee!e!ffd0e!d1e$fd6Z'd d d ddddddddddd e e dddddddd d d d d ddd fd7ee!ee!e!fee!e!e!ffd8e$d9e$d ee"deee#e#fdeee#e#fd!e#d"e#d#ee#ee#dffd$ee#d&e#d'e#d(ee"de"dee#dfdee#dfd)e#d*e"d+e!d,e!d3ee#d4ee"d5ee!d:e$de$de$d-e$d.e$d/ee!ee!e!ffd0e!d1e$f>d;Z(y)| t=j>|!gz } |dkDr| tA||||d"#gz } |r| tC|$gz } |r>t jD|t jD|t jD| fSt jD||z| zScc}w)%a ImageNet-oriented image transforms for training. Args: img_size: Target image size. train_crop_mode: Training random crop mode ('rrc', 'rkrc', 'rkrr'). scale: Random resize scale range (crop area, < 1.0 => zoom in). ratio: Random aspect ratio range (crop ratio for RRC, ratio adjustment factor for RKR). hflip: Horizontal flip probability. vflip: Vertical flip probability. color_jitter: Random color jitter component factors (brightness, contrast, saturation, hue). Scalar is applied as (scalar,) * 3 (no hue). color_jitter_prob: Apply color jitter with this probability if not None (for SimlCLR-like aug). force_color_jitter: Force color jitter where it is normally disabled (ie with RandAugment on). grayscale_prob: Probability of converting image to grayscale (for SimCLR-like aug). gaussian_blur_prob: Probability of applying gaussian blur (for SimCLR-like aug). auto_augment: Auto augment configuration string (see auto_augment.py). interpolation: Image interpolation mode. mean: Image normalization mean. std: Image normalization standard deviation. re_prob: Random erasing probability. re_mode: Random erasing fill mode. re_count: Number of random erasing regions. re_num_splits: Control split of random erasing across batch size. use_prefetcher: Prefetcher enabled. Do not convert image to tensor or normalize. normalize: Normalize tensor output w/ provided mean/std (if prefetcher not used). separate: Output transforms in 3-stage tuple. naflex: Enable NaFlex mode, sequence constrained patch output patch_size: Patch size for NaFlex mode. max_seq_len: Max sequence length for NaFlex mode. Returns: If separate==True, the transforms are returned as a tuple of 3 separate transforms for use in a mixing dataset that passes * all data through the first (primary) transform, called the 'clean' data * a portion of the data through the secondary transform * normalizes and converts the branches above with the third, final transform rrc>rJrkrcrkrr)g{Gz??)g?gUUUUUU?)rFrGr5r6r)rKrL)g?rM)g?grq?r0T)rrandom_scale_probrandom_scale_rangerandom_scale_arearandom_aspect_probrandom_aspect_rangerKreflect) padding_mode)r5r6rr1)pF3ag?)translate_constimg_meanr#rrandaugmixg333333? translate_pct)r]) kernel_sizer%cpu)mode max_count num_splitsdevicerF)#tuplerrrrrrRandomHorizontalFlipRandomVerticalFlip isinstancestrlistmindictintroundr startswithr r r lenfloat RandomApply ColorJitterRandomGrayscale GaussianBlurrrr(r)r*rrr+)!rr5r6r7r8r9r:r;r<r=r>r?rrrr@rArBrCr r!rDrErFrGrH primary_tfl secondary_tfldisable_color_jitter img_size_minx aa_params final_tfls! r-transforms_imagenet_trainrAsXB&.O 5 55 5K e*{+e1123!#'     . .%.;/E%.;/E "/&)',&*'*(-#f, yAyA  K%.;/E%5#56E 1"/  K rz 77%@AA  rz 55>?? M ,,,,%7$N$,:NO h .x=L#L t 34dCCU37^4CD  ]h6):=)IIo &  " "6 * 4\9MN NM  $ $X .),Io & 7 iPQ QM 4\9MN NM(< lT5M 2|$. ..",/1A5L  ( &&".. =(( M j44lCD DM*44~FGG   " "++;$%   I&()) &())  O  \\$'LL%    R<  &,  Ih*566 !!+. 0B0B=0QS]SeSefoSppp!!+ "= "IJJYDs*Ocrop_pct crop_modecrop_border_pixelsc 4|xst}t|ttfr@t |dk(sJt|D cgc]} t j | |z c} }nt j ||z }||f}g}|r|t|gz }| r|t| | |gz }n|dk(r9|tj|t|tj|gz }n|dk(r:|Dcgc]}td|z}}|t||dt|| gz }n[|d |d k(r(|tj|d t|gz }n|t|gz }|tj|gz }|r|t!gz }n[|s|t!gz }nJ|t#tj$t'j(|t'j(| gz }| r|t+| gz }tj,|Scc} wcc}w)a ImageNet-oriented image transform for evaluation and inference. Args: img_size: Target image size. crop_pct: Crop percentage. Defaults to 0.875 when None. crop_mode: Crop mode. One of ['squash', 'border', 'center']. Defaults to 'center' when None. crop_border_pixels: Trim a border of specified # pixels around edge of original image. interpolation: Image interpolation mode. mean: Image normalization mean. std: Image normalization standard deviation. use_prefetcher: Prefetcher enabled. Do not convert image to tensor or normalize. normalize: Normalize tensor output w/ provided mean/std (if prefetcher not used). naflex: Enable NaFlex mode, sequence constrained patch output patch_size: Patch size for NaFlex mode. max_seq_len: Max sequence length for NaFlex mode. patchify: Patchify the output instead of relying on prefetcher Returns: Composed transform pipeline )rFrGrsquashr$borderrWrM)rlongest)fillrr3r%rf)r rjrgrlrrmathfloorrrrr&r r'rprrrrr(r)r*rr+)rrrrrrrr r!rErFrGrHr| scale_sizer,vrs r-transforms_imagenet_evalrs7F++H(UDM*8}!!!hGDJJq8|4GH ZZ8 34  *- C -.//   !#'      !!* zoom in). ratio: Random aspect ratio range (crop ratio for RRC, ratio adjustment factor for RKR). hflip: Horizontal flip probability. vflip: Vertical flip probability. color_jitter: Random color jitter component factors (brightness, contrast, saturation, hue). Scalar is applied as (scalar,) * 3 (no hue). color_jitter_prob: Apply color jitter with this probability if not None (for SimlCLR-like aug). grayscale_prob: Probability of converting image to grayscale (for SimCLR-like aug). gaussian_blur_prob: Probability of applying gaussian blur (for SimCLR-like aug). auto_augment: Auto augment configuration string (see auto_augment.py). interpolation: Image interpolation mode. mean: Image normalization mean. std: Image normalization standard deviation. re_prob: Random erasing probability. re_mode: Random erasing fill mode. re_count: Number of random erasing regions. re_num_splits: Control split of random erasing across batch size. crop_pct: Inference crop percentage (output size / resize size). crop_mode: Inference crop mode. One of ['squash', 'border', 'center']. Defaults to 'center' when None. crop_border_pixels: Inference crop border of specified # pixels around edge of original image. tf_preprocessing: Use TF 1.0 inference preprocessing for testing model ports use_prefetcher: Pre-fetcher enabled. Do not convert image to tensor or normalize. normalize: Normalization tensor output w/ provided mean/std (if prefetcher not used). separate: Output transforms in 3-stage tuple. Returns: Composed transforms or tuple thereof Nz6Separate transforms not supported for TF preprocessingr)TfPreprocessTransform)rsizerz-Cannot perform split augmentation with no_aug)rrrr r!r7r5r6r8r9r:r;r=r>r?rrrr@rArBrCr r!rDrErFrGrHz>Separate transforms not supported for validation preprocessing) rrrrrrr r!rErFrGrH)rjrgrltimm.data.tf_preprocessingrr.rr)"rrrr7r5r6r8r9r:r;r=r>r?rrrr@rArBrCrrrrr r!rDrErFrGrHrr transforms" r-create_transformrysH*udm,bc?NUUU|D)#' | q 6 P!P P<.+-# Il ]1 /     *#4 .$6*,   !""#$,%& .'($)*"+,-.&/0(12"3IZ # a!a a<0+!##5-#%'!I r/))__doc__rtypingrrrr) torchvisionrtimm.data.constantsrrr timm.data.auto_augmentr r r timm.data.transformsr rrrrrrrrtimm.data.naflex_transformsrrrtimm.data.random_erasingrrorkrsboolr.rrrr/r-rs  )) "]]llccc__214'"7!5$ *#U38_,-*#*#E3J*#5#:  *#  *#  *#\14/3/3)-8;-1#( "$&&*%"7!5$245LKU38_,-LKeUl+,LKeUl+,LK"# LK  LK  LKE5#445LK$E?LK!LKLK"LKsmLKLKE3JLK5#: LK !LK"#LK$%LK&'LK()LK*+LK,-LK./LK0#uS#X./1LK23LK45LK`14$(#',0'"7!5$24f#U38_,-f#5/f#C=f#%SM f#  f# E3J f#5#: f#f#f#f##uS#X./f#f#f#TIL!)-/3/38;-1 "$&&*'"7!5$(#',0!&$24?J#uS#Xc3m0DDEJJJ"# J eUl+, J eUl+, JJJE5#445J$E?JJ"JsmJJE3JJ 5#: !J"#J$%J&'J()J*5/+J,C=-J.%SM/J01J23J45J67J89J:#uS#X./;J<=J>?Jr/