No public description

PiperOrigin-RevId: 553017713

official/vision/dataloaders/retinanet_input.py

@@ -177,7 +177,7 @@ class Parser(parser.Parser):
                                                  image_info[1, :], offset)
     return image, boxes, image_info
 
-  def _parse_train_data(self, data, anchor_labeler=None):
+  def _parse_train_data(self, data, anchor_labeler=None, input_anchor=None):
     """Parses data for training and evaluation."""
     classes = data['groundtruth_classes']
     boxes = data['groundtruth_boxes']
@@ -251,12 +251,15 @@ class Parser(parser.Parser):
       attributes[k] = tf.gather(v, indices)
 
     # Assigns anchors.
-    input_anchor = anchor.build_anchor_generator(
-        min_level=self._min_level,
-        max_level=self._max_level,
-        num_scales=self._num_scales,
-        aspect_ratios=self._aspect_ratios,
-        anchor_size=self._anchor_size)
+    if input_anchor is None:
+      input_anchor = anchor.build_anchor_generator(
+          min_level=self._min_level,
+          max_level=self._max_level,
+          num_scales=self._num_scales,
+          aspect_ratios=self._aspect_ratios,
+          anchor_size=self._anchor_size,
+      )
+
     anchor_boxes = input_anchor(image_size=(image_height, image_width))
     if anchor_labeler is None:
       anchor_labeler = anchor.AnchorLabeler(
@@ -284,7 +287,7 @@ class Parser(parser.Parser):
       labels['attribute_targets'] = att_targets
     return image, labels
 
-  def _parse_eval_data(self, data, anchor_labeler=None):
+  def _parse_eval_data(self, data, anchor_labeler=None, input_anchor=None):
     """Parses data for training and evaluation."""
 
     classes = data['groundtruth_classes']
@@ -326,12 +329,15 @@ class Parser(parser.Parser):
       attributes[k] = tf.gather(v, indices)
 
     # Assigns anchors.
-    input_anchor = anchor.build_anchor_generator(
-        min_level=self._min_level,
-        max_level=self._max_level,
-        num_scales=self._num_scales,
-        aspect_ratios=self._aspect_ratios,
-        anchor_size=self._anchor_size)
+    if input_anchor is None:
+      input_anchor = anchor.build_anchor_generator(
+          min_level=self._min_level,
+          max_level=self._max_level,
+          num_scales=self._num_scales,
+          aspect_ratios=self._aspect_ratios,
+          anchor_size=self._anchor_size,
+      )
+
     anchor_boxes = input_anchor(image_size=(image_height, image_width))
     if anchor_labeler is None:
       anchor_labeler = anchor.AnchorLabeler(

official/vision/ops/anchor.py

@@ -15,6 +15,7 @@
 """Anchor box and labeler definition."""
 
 import collections
+import math
 from typing import Dict, Optional, Tuple
 
 # Import libraries
@@ -78,9 +79,10 @@ class Anchor(object):
     boxes_all = []
     for level in range(self.min_level, self.max_level + 1):
       boxes_l = []
+      feat_size = math.ceil(self.image_size[0] / 2**level)
+      stride = tf.cast(self.image_size[0] / feat_size, tf.float32)
       for scale in range(self.num_scales):
         for aspect_ratio in self.aspect_ratios:
-          stride = 2**level
           intermidate_scale = 2 ** (scale / float(self.num_scales))
           base_anchor_size = self.anchor_size * stride * intermidate_scale
           aspect_x = aspect_ratio**0.5