workaround for iteration valid sample check

Author: HiKapok

preprocessing/ssd_preprocessing.py

@@ -315,6 +315,36 @@ def ssd_random_expand(image, bboxes, ratio=2., name=None):
 
     return big_canvas, absolute_bboxes / tf.cast(tf.stack([canvas_height, canvas_width, canvas_height, canvas_width]), bboxes.dtype)
 
+# def ssd_random_sample_patch_wrapper(image, labels, bboxes):
+#   with tf.name_scope('ssd_random_sample_patch_wrapper'):
+#     orgi_image, orgi_labels, orgi_bboxes = image, labels, bboxes
+#     def check_bboxes(bboxes):
+#       areas = (bboxes[:, 3] - bboxes[:, 1]) * (bboxes[:, 2] - bboxes[:, 0])
+#       return tf.logical_and(tf.logical_and(areas < 0.9, areas > 0.001),
+#                             tf.logical_and((bboxes[:, 3] - bboxes[:, 1]) > 0.025, (bboxes[:, 2] - bboxes[:, 0]) > 0.025))
+
+#     index = 0
+#     max_attempt = 3
+#     def condition(index, image, labels, bboxes):
+#       return tf.logical_or(tf.logical_and(tf.reduce_sum(tf.cast(check_bboxes(bboxes), tf.int64)) < 1, tf.less(index, max_attempt)), tf.less(index, 1))
+
+#     def body(index, image, labels, bboxes):
+#       image, bboxes = tf.cond(tf.random_uniform([], minval=0., maxval=1., dtype=tf.float32) < 0.5,
+#                       lambda: (image, bboxes),
+#                       lambda: ssd_random_expand(image, bboxes, tf.random_uniform([1], minval=1.1, maxval=4., dtype=tf.float32)[0]))
+#       # Distort image and bounding boxes.
+#       random_sample_image, labels, bboxes = ssd_random_sample_patch(image, labels, bboxes, ratio_list=[-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.])
+#       random_sample_image.set_shape([None, None, 3])
+#       return index+1, random_sample_image, labels, bboxes
+
+#     [index, image, labels, bboxes] = tf.while_loop(condition, body, [index, orgi_image, orgi_labels, orgi_bboxes], parallel_iterations=4, back_prop=False, swap_memory=True)
+
+#     valid_mask = check_bboxes(bboxes)
+#     labels, bboxes = tf.boolean_mask(labels, valid_mask), tf.boolean_mask(bboxes, valid_mask)
+#     return tf.cond(tf.less(index, max_attempt),
+#                 lambda : (image, labels, bboxes),
+#                 lambda : (orgi_image, orgi_labels, orgi_bboxes))
+
 def ssd_random_sample_patch_wrapper(image, labels, bboxes):
   with tf.name_scope('ssd_random_sample_patch_wrapper'):
     orgi_image, orgi_labels, orgi_bboxes = image, labels, bboxes
@@ -324,20 +354,20 @@ def check_bboxes(bboxes):
                             tf.logical_and((bboxes[:, 3] - bboxes[:, 1]) > 0.025, (bboxes[:, 2] - bboxes[:, 0]) > 0.025))
 
     index = 0
-    max_attempt = 1
-    def condition(index, image, labels, bboxes):
+    max_attempt = 3
+    def condition(index, image, labels, bboxes, orgi_image, orgi_labels, orgi_bboxes):
       return tf.logical_or(tf.logical_and(tf.reduce_sum(tf.cast(check_bboxes(bboxes), tf.int64)) < 1, tf.less(index, max_attempt)), tf.less(index, 1))
 
-    def body(index, image, labels, bboxes):
+    def body(index, image, labels, bboxes, orgi_image, orgi_labels, orgi_bboxes):
       image, bboxes = tf.cond(tf.random_uniform([], minval=0., maxval=1., dtype=tf.float32) < 0.5,
-                      lambda: (image, bboxes),
-                      lambda: ssd_random_expand(image, bboxes, tf.random_uniform([1], minval=1.1, maxval=4., dtype=tf.float32)[0]))
+                      lambda: (orgi_image, orgi_bboxes),
+                      lambda: ssd_random_expand(orgi_image, orgi_bboxes, tf.random_uniform([1], minval=1.1, maxval=4., dtype=tf.float32)[0]))
       # Distort image and bounding boxes.
-      random_sample_image, labels, bboxes = ssd_random_sample_patch(image, labels, bboxes, ratio_list=[-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.])
+      random_sample_image, labels, bboxes = ssd_random_sample_patch(image, orgi_labels, bboxes, ratio_list=[-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.])
       random_sample_image.set_shape([None, None, 3])
-      return index+1, random_sample_image, labels, bboxes
+      return index+1, random_sample_image, labels, bboxes, orgi_image, orgi_labels, orgi_bboxes
 
-    [index, image, labels, bboxes] = tf.while_loop(condition, body, [index, orgi_image, orgi_labels, orgi_bboxes], parallel_iterations=4, back_prop=False, swap_memory=True)
+    [index, image, labels, bboxes, orgi_image, orgi_labels, orgi_bboxes] = tf.while_loop(condition, body, [index,  image, labels, bboxes, orgi_image, orgi_labels, orgi_bboxes], parallel_iterations=4, back_prop=False, swap_memory=True)
 
     valid_mask = check_bboxes(bboxes)
     labels, bboxes = tf.boolean_mask(labels, valid_mask), tf.boolean_mask(bboxes, valid_mask)