merge prior_box and multi_box

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

@@ -109,16 +109,19 @@ 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):
+        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)
@@ -127,46 +130,19 @@ class TestPriorBox(unittest.TestCase):
         conv4 = fluid.layers.conv2d(conv3, 3, 3, 2)
         conv5 = fluid.layers.conv2d(conv4, 3, 3, 2)
 
-        box, var = layers.prior_box(
+        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
-
-
-class TestMultiBoxHead(unittest.TestCase):
-    def test_prior_box(self):
-        data_shape = [3, 224, 224]
-        mbox_locs, mbox_confs = self.multi_box_output(data_shape)
-
-        for loc, conf in zip(mbox_locs, mbox_confs):
-            assert loc.shape[1:3] == conf.shape[1:3]
-
-    def multi_box_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 = detection.multi_box_head(
-            inputs=[conv1, conv2, conv3, conv4, conv5, conv5],
-            num_classes=21,
-            min_ratio=20,
-            max_ratio=90,
-            aspect_ratios=[[2.], [2., 3.], [2., 3.], [2., 3.], [2.], [2.]],
-            base_size=300,
-            flip=True)
-        return mbox_locs, mbox_confs
+        return mbox_locs, mbox_confs, box, var
 
 
 if __name__ == '__main__':