2020
2121from __future__ import absolute_import
2222from __future__ import division
23+ # gtype import
2324from __future__ import print_function
2425
2526import inspect
2627import math
2728from absl import logging
29+ import numpy as np
2830import tensorflow .compat .v1 as tf
2931import tensorflow_probability as tfp
32+
3033import hparams_config
31- import numpy as np
3234
3335try :
3436 # addon image_ops are simpler, but they have some issues on GPU and TPU.
@@ -1547,7 +1549,13 @@ def select_and_apply_random_policy(policies, image, bboxes):
15471549 lambda : (image , bboxes ))
15481550 return (image , bboxes )
15491551
1550- def select_and_apply_random_policy_augmix (policies , image , bboxes , mixture_width = 3 , mixture_depth = - 1 , alpha = 1 ):
1552+
1553+ def select_and_apply_random_policy_augmix (policies ,
1554+ image ,
1555+ bboxes ,
1556+ mixture_width = 3 ,
1557+ mixture_depth = - 1 ,
1558+ alpha = 1 ):
15511559 """Select a random policy from `policies` and apply it to `image`."""
15521560 policy_to_select = tf .random_uniform ([], maxval = len (policies ), dtype = tf .int32 )
15531561 # Note that using tf.case instead of tf.conds would result in significantly
@@ -1562,12 +1570,13 @@ def select_and_apply_random_policy_augmix(policies, image, bboxes, mixture_width
15621570 for (i , policy ) in enumerate (policies ):
15631571 aug_image , bboxes = tf .cond (
15641572 tf .equal (i , policy_to_select ),
1565- lambda selected_policy = policy : selected_policy ( aug_image , bboxes ),
1566- lambda : (aug_image , bboxes ))
1573+ lambda policy_fn = policy , img = aug_image : policy_fn ( img , bboxes ),
1574+ lambda img = aug_image : (img , bboxes ))
15671575 mix += ws [j ] * tf .cast (aug_image , tf .float32 )
15681576 mixed = tf .cast ((1 - m ) * tf .cast (image , tf .float32 ) + m * mix , tf .uint8 )
15691577 return (mixed , bboxes )
15701578
1579+
15711580def build_and_apply_nas_policy (policies , image , bboxes ,
15721581 augmentation_hparams , use_augmix = False ,
15731582 mixture_width = 3 , mixture_depth = - 1 , alpha = 1 ):
@@ -1583,9 +1592,9 @@ def build_and_apply_nas_policy(policies, image, bboxes,
15831592 normalized between [0, 1].
15841593 augmentation_hparams: Hparams associated with the NAS learned policy.
15851594 use_augmix: whether use augmix[https://arxiv.org/pdf/1912.02781.pdf]
1586- width : Width of augmentation chain
1587- depth : Depth of augmentation chain. -1 enables stochastic depth uniformly
1588- from [1, 3]
1595+ mixture_width : Width of augmentation chain
1596+ mixture_depth : Depth of augmentation chain. -1 enables stochastic depth
1597+ uniformly from [1, 3].
15891598 alpha: Probability coefficient for Beta and Dirichlet distributions.
15901599
15911600 Returns:
@@ -1631,8 +1640,13 @@ def final_policy(image_, bboxes_):
16311640 return (augmented_images , augmented_bboxes )
16321641
16331642
1634- def distort_image_with_autoaugment (image , bboxes , augmentation_name , use_augmix = False ,
1635- mixture_width = 3 , mixture_depth = - 1 , alpha = 1 ):
1643+ def distort_image_with_autoaugment (image ,
1644+ bboxes ,
1645+ augmentation_name ,
1646+ use_augmix = False ,
1647+ mixture_width = 3 ,
1648+ mixture_depth = - 1 ,
1649+ alpha = 1 ):
16361650 """Applies the AutoAugment policy to `image` and `bboxes`.
16371651
16381652 Args:
@@ -1646,6 +1660,11 @@ def distort_image_with_autoaugment(image, bboxes, augmentation_name, use_augmix=
16461660 found on the COCO dataset that have slight variation in what operations
16471661 were used during the search procedure along with how many operations are
16481662 applied in parallel to a single image (2 vs 3).
1663+ use_augmix: whether use augmix[https://arxiv.org/pdf/1912.02781.pdf]
1664+ mixture_width: Width of augmentation chain
1665+ mixture_depth: Depth of augmentation chain. -1 enables stochastic depth
1666+ uniformly from [1, 3].
1667+ alpha: Probability coefficient for Beta and Dirichlet distributions.
16491668
16501669 Returns:
16511670 A tuple containing the augmented versions of `image` and `bboxes`.
@@ -1668,4 +1687,5 @@ def distort_image_with_autoaugment(image, bboxes, augmentation_name, use_augmix=
16681687
16691688 with tf .device ('/cpu:0' ):
16701689 return build_and_apply_nas_policy (policy , image , bboxes ,
1671- augmentation_hparams , use_augmix , mixture_width , mixture_depth , alpha )
1690+ augmentation_hparams , use_augmix ,
1691+ mixture_width , mixture_depth , alpha )
0 commit comments