Fix testing in pose model for COCO (#3682)

PaddleCV/human_pose_estimation/README.md

@@ -108,9 +108,12 @@ We also support to apply pre-trained models on customized images.
 Put the images into the folder `test` under the directory root of this repo. Then run
 
 ```bash
+# default is MPII dataset
 python test.py --checkpoint checkpoints/pose-resnet-50-384x384-mpii
 ```
 
+`python test.py --help` for more usage.
+
 If there are multiple persons in images, detectors such as [Faster R-CNN](https://github.com/PaddlePaddle/models/tree/develop/fluid/PaddleCV/rcnn), [SSD](https://github.com/PaddlePaddle/models/tree/develop/fluid/PaddleCV/object_detection) or others should be used first to crop them out. Because the simple baseline for human pose estimation is a top-down method.
 
 ## Reference

PaddleCV/human_pose_estimation/README_cn.md

@@ -107,9 +107,12 @@ python train.py --dataset 'mpii'
 将测试图片放入根目录下的'test'文件夹中，执行
 
 ```bash
+# 默认是MPII数据集
 python test.py --checkpoint 'checkpoints/pose-resnet-50-384x384-mpii'
 ```
 
+`python test.py --help`获取更多的用法。
+
 ## 引用
 
 - Simple Baselines for Human Pose Estimation and Tracking in PyTorch [`code`](https://github.com/Microsoft/human-pose-estimation.pytorch#data-preparation)

PaddleCV/human_pose_estimation/test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 ##############################################################################
-
 """Functions for inference."""
 
 import sys
@@ -46,23 +45,27 @@ def print_immediately(s):
 
 
 def test(args):
-    import lib.mpii_reader as reader
     if args.dataset == 'coco':
+        import lib.coco_reader as reader
         IMAGE_SIZE = [288, 384]
-        FLIP_PAIRS = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]]
+        FLIP_PAIRS = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12],
+                      [13, 14], [15, 16]]
         args.kp_dim = 17
     elif args.dataset == 'mpii':
+        import lib.mpii_reader as reader
         IMAGE_SIZE = [384, 384]
         FLIP_PAIRS = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
         args.kp_dim = 16
     else:
-        raise ValueError('The dataset {} is not supported yet.'.format(args.dataset))
+        raise ValueError('The dataset {} is not supported yet.'.format(
+            args.dataset))
 
     print_arguments(args)
 
     # Image and target
-    image = layers.data(name='image', shape=[3, IMAGE_SIZE[1], IMAGE_SIZE[0]], dtype='float32')
-    file_id = layers.data(name='file_id', shape=[1,], dtype='int')
+    image = layers.data(
+        name='image', shape=[3, IMAGE_SIZE[1], IMAGE_SIZE[0]], dtype='float32')
+    file_id = layers.data(name='file_id', shape=[1, ], dtype='int')
 
     # Build model
     model = pose_resnet.ResNet(layers=50, kps_num=args.kp_dim, test_mode=True)
@@ -82,9 +85,9 @@ def test(args):
     feeder = fluid.DataFeeder(place=place, feed_list=[image, file_id])
 
     test_exe = fluid.ParallelExecutor(
-            use_cuda=True if args.use_gpu else False,
-            main_program=fluid.default_main_program().clone(for_test=True),
-            loss_name=None)
+        use_cuda=True if args.use_gpu else False,
+        main_program=fluid.default_main_program().clone(for_test=True),
+        loss_name=None)
 
     fetch_list = [image.name, output.name]
 
@@ -96,22 +99,18 @@ def test(args):
         for i in range(num_images):
             file_ids.append(data[i][1])
 
-        input_image, out_heatmaps  = test_exe.run(
-                fetch_list=fetch_list,
-                feed=feeder.feed(data))
+        input_image, out_heatmaps = test_exe.run(fetch_list=fetch_list,
+                                                 feed=feeder.feed(data))
 
         if args.flip_test:
             # Flip all the images in a same batch
             data_fliped = []
             for i in range(num_images):
-                data_fliped.append((
-                            data[i][0][:, :, ::-1],
-                            data[i][1]))
+                data_fliped.append((data[i][0][:, :, ::-1], data[i][1]))
 
             # Inference again
-            _, output_flipped = test_exe.run(
-                    fetch_list=fetch_list,
-                    feed=feeder.feed(data_fliped))
+            _, output_flipped = test_exe.run(fetch_list=fetch_list,
+                                             feed=feeder.feed(data_fliped))
 
             # Flip back
             output_flipped = flip_back(output_flipped, FLIP_PAIRS)
@@ -123,7 +122,8 @@ def test(args):
 
             # Aggregate
             out_heatmaps = (out_heatmaps + output_flipped) * 0.5
-            save_predict_results(input_image, out_heatmaps, file_ids, fold_name='results')
+            save_predict_results(
+                input_image, out_heatmaps, file_ids, fold_name='results')
 
 
 if __name__ == '__main__':