Merge pull request #8382 from chengduoZH/feature/multiBoxHead

Add MultiBox API

python/paddle/v2/fluid/tests/test_detection.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from __future__ import print_function
+import paddle.v2.fluid as fluid
 import paddle.v2.fluid.layers as layers
 from paddle.v2.fluid.framework import Program, program_guard
 import unittest
@@ -108,60 +109,40 @@ class TestDetection(unittest.TestCase):
         print(str(program))
 
 
-class TestPriorBox(unittest.TestCase):
-    def test_prior_box(self):
+class TestMultiBoxHead(unittest.TestCase):
+    def test_multi_box_head(self):
         data_shape = [3, 224, 224]
-        box, var = self.prior_box_output(data_shape)
+        mbox_locs, mbox_confs, box, var = self.multi_box_head_output(data_shape)
 
         assert len(box.shape) == 2
         assert box.shape == var.shape
         assert box.shape[1] == 4
 
-    def prior_box_output(self, data_shape):
-        images = layers.data(name='pixel', shape=data_shape, dtype='float32')
-        conv1 = layers.conv2d(
-            input=images,
-            num_filters=3,
-            filter_size=3,
-            stride=2,
-            use_cudnn=False)
-        conv2 = layers.conv2d(
-            input=conv1,
-            num_filters=3,
-            filter_size=3,
-            stride=2,
-            use_cudnn=False)
-        conv3 = layers.conv2d(
-            input=conv2,
-            num_filters=3,
-            filter_size=3,
-            stride=2,
-            use_cudnn=False)
-        conv4 = layers.conv2d(
-            input=conv3,
-            num_filters=3,
-            filter_size=3,
-            stride=2,
-            use_cudnn=False)
-        conv5 = layers.conv2d(
-            input=conv4,
-            num_filters=3,
-            filter_size=3,
-            stride=2,
-            use_cudnn=False)
-
-        box, var = layers.prior_box(
+        for loc, conf in zip(mbox_locs, mbox_confs):
+            assert loc.shape[1:3] == conf.shape[1:3]
+
+    def multi_box_head_output(self, data_shape):
+        images = fluid.layers.data(
+            name='pixel', shape=data_shape, dtype='float32')
+        conv1 = fluid.layers.conv2d(images, 3, 3, 2)
+        conv2 = fluid.layers.conv2d(conv1, 3, 3, 2)
+        conv3 = fluid.layers.conv2d(conv2, 3, 3, 2)
+        conv4 = fluid.layers.conv2d(conv3, 3, 3, 2)
+        conv5 = fluid.layers.conv2d(conv4, 3, 3, 2)
+
+        mbox_locs, mbox_confs, box, var = layers.multi_box_head(
             inputs=[conv1, conv2, conv3, conv4, conv5, conv5],
             image=images,
+            num_classes=21,
             min_ratio=20,
             max_ratio=90,
-            # steps=[8, 16, 32, 64, 100, 300],
             aspect_ratios=[[2.], [2., 3.], [2., 3.], [2., 3.], [2.], [2.]],
             base_size=300,
             offset=0.5,
             flip=True,
             clip=True)
-        return box, var
+
+        return mbox_locs, mbox_confs, box, var
 
 
 if __name__ == '__main__':