Add docs for detection generator

official/vision/beta/projects/yolo/modeling/layers/detection_generator.py

@@ -40,9 +40,6 @@ class YoloLayer(ks.Model):
     """
     parameters for the loss functions used at each detection head output
 
-scale_anchors: `int` for how much to scale this level to get the orginal
-  input shape
-
     Args:
       masks: `List[int]` for the output level that this specific model output
         level.
@@ -59,18 +56,23 @@ scale_anchors: `int` for how much to scale this level to get the orginal
       max_delta: gradient clipping to apply to the box loss.
       loss_type: `str` for the typeof iou loss to use with in {ciou, diou,
         giou, iou}.
-
+      use_tie_breaker: TODO unused?
       iou_normalizer: `float` for how much to scale the loss on the IOU or the
         boxes.
       cls_normalizer: `float` for how much to scale the loss on the classes.
       obj_normalizer: `float` for how much to scale loss on the detection map.
-      objectness_smooth: `float` for how much to smooth the loss on the
-        detection map.
       use_scaled_loss: `bool` for whether to use the scaled loss
         or the traditional loss.
-
+      darknet: `bool` for whether to use the DarkNet or PyTorch loss function
+        implementation.
+      pre_nms_points: `int` number of top candidate detections per class before
+        NMS.
       label_smoothing: `float` for how much to smooth the loss on the classes.
-      new_cords: `bool` for which scaling type to use.
+      max_boxes: `int` for the maximum number of boxes retained over all
+        classes.
+      new_cords: `bool` for using the ScaledYOLOv4 coordinates.
+      path_scale: `dict` for the size of the input tensors. Defaults to
+        precalulated values from the `mask`.
       scale_xy: dictionary `float` values inidcating how far each pixel can see
         outside of its containment of 1.0. a value of 1.2 indicates there is a
         20% extended radius around each pixel that this specific pixel can
@@ -78,11 +80,9 @@ scale_anchors: `int` for how much to scale this level to get the orginal
         to 1 + value/2, this value is set in the yolo filter, and resused here.
         there should be one value for scale_xy for each level from min_level to
         max_level.
-      nms_type: "greedy",
-      nms_thresh: 0.6,
-      iou_thresh: 0.213,
-      name=None,
-
+      nms_type: `str` for which non max suppression to use.
+      objectness_smooth: `float` for how much to smooth the loss on the
+        detection map.
 
     Return:
       loss: `float` for the actual loss.

official/vision/beta/projects/yolo/modeling/layers/detection_generator_test.py

+# Lint as: python3
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for yolo."""
+
+# Import libraries
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+# from official.vision.beta.projects.yolo.modeling.backbones import darknet
+from official.vision.beta.projects.yolo.modeling.layers import detection_generator as dg
+
+
+class YoloDecoderTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (True),
+      (False),
+  )
+  def test_network_creation(self, nms):
+    """Test creation of ResNet family models."""
+    tf.keras.backend.set_image_data_format('channels_last')
+    input_shape = {
+        '3': [1, 52, 52, 255],
+        '4': [1, 26, 26, 255],
+        '5': [1, 13, 13, 255]
+    }
+    classes = 80
+    masks = {'3': [0, 1, 2], '4': [3, 4, 5], '5': [6, 7, 8]}
+    anchors = [[12.0, 19.0], [31.0, 46.0], [96.0, 54.0], [46.0, 114.0],
+               [133.0, 127.0], [79.0, 225.0], [301.0, 150.0], [172.0, 286.0],
+               [348.0, 340.0]]
+    layer = dg.YoloLayer(masks, anchors, classes, max_boxes=10)
+
+    inputs = {}
+    for key in input_shape.keys():
+      inputs[key] = tf.ones(input_shape[key], dtype=tf.float32)
+
+    endpoints = layer(inputs)
+
+    boxes = endpoints['bbox']
+    classes = endpoints['classes']
+
+    self.assertAllEqual(boxes.shape.as_list(), [1, 10, 4])
+    self.assertAllEqual(classes.shape.as_list(), [1, 10])
+
+
+if __name__ == '__main__':
+  from yolo.utils.run_utils import prep_gpu
+  prep_gpu()
+  tf.test.main()