update pp-tinypose (#1974)

* update pp-tinypose

* update version

* add clean func

* update README

modules/image/keypoint_detection/pp-tinypose/README.md

@@ -59,7 +59,7 @@
     ```
   - 通过命令行方式实现关键点检测模型的调用，更多请见 [PaddleHub命令行指令](../../../../docs/docs_ch/tutorial/cmd_usage.rst)
 
-- ### 2、代码示例
+- ### 2、预测代码示例
 
   - ```python
     import paddlehub as hub
@@ -131,6 +131,10 @@
 
   初始发布
 
+* 1.1.0
+
+  修复使用 ndarray 输入时无法保存可视化图片的问题
+
   - ```shell
-    $ hub install pp-tinypose==1.0.0
+    $ hub install pp-tinypose==1.1.0
     ```

modules/image/keypoint_detection/pp-tinypose/__init__.py

-# Copyright (c) 2022 PaddlePaddle 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.
-import glob
-import os
-import sys
-
-CUR_DIR = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(CUR_DIR)

modules/image/keypoint_detection/pp-tinypose/det_keypoint_unite_infer.py

@@ -11,21 +11,9 @@
 # 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.
-import json
-import math
-import os
-
-import cv2
 import numpy as np
-import paddle
-import yaml
-from infer import Detector
-from infer import PredictConfig
-from keypoint_infer import KeyPointDetector
-from keypoint_infer import PredictConfig_KeyPoint
-from keypoint_postprocess import translate_to_ori_images
-from preprocess import decode_image
-from visualize import visualize_pose
+
+from .keypoint_postprocess import translate_to_ori_images
 
 KEYPOINT_SUPPORT_MODELS = {'HigherHRNet': 'keypoint_bottomup', 'HRNet': 'keypoint_topdown'}
 

modules/image/keypoint_detection/pp-tinypose/infer.py

@@ -15,29 +15,27 @@ import glob
 import json
 import math
 import os
-import sys
-from functools import reduce
 from pathlib import Path
 
 import cv2
 import numpy as np
-import paddle
 import yaml
-from keypoint_preprocess import EvalAffine
-from keypoint_preprocess import expand_crop
-from keypoint_preprocess import TopDownEvalAffine
 from paddle.inference import Config
 from paddle.inference import create_predictor
-from preprocess import decode_image
-from preprocess import LetterBoxResize
-from preprocess import NormalizeImage
-from preprocess import Pad
-from preprocess import PadStride
-from preprocess import Permute
-from preprocess import preprocess
-from preprocess import Resize
-from preprocess import WarpAffine
-from visualize import visualize_box
+
+from .preprocess import decode_image
+from .preprocess import LetterBoxResize
+from .preprocess import NormalizeImage
+from .preprocess import Pad
+from .preprocess import PadStride
+from .preprocess import Permute
+from .preprocess import preprocess
+from .preprocess import Resize
+from .preprocess import WarpAffine
+from .keypoint_preprocess import EvalAffine
+from .keypoint_preprocess import expand_crop
+from .keypoint_preprocess import TopDownEvalAffine
+from .visualize import visualize_box
 
 # Global dictionary
 SUPPORT_MODELS = {

modules/image/keypoint_detection/pp-tinypose/keypoint_infer.py

@@ -11,29 +11,17 @@
 # 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.
-import glob
 import math
 import os
-import sys
-import time
-from functools import reduce
 
 import cv2
 import numpy as np
-import paddle
 import yaml
-from PIL import Image
-# add deploy path of PadleDetection to sys.path
-parent_path = os.path.abspath(os.path.join(__file__, *(['..'])))
-sys.path.insert(0, parent_path)
 
-from preprocess import preprocess, NormalizeImage, Permute
-from keypoint_preprocess import EvalAffine, TopDownEvalAffine, expand_crop
-from keypoint_postprocess import HRNetPostProcess
-from visualize import visualize_pose
-from paddle.inference import Config
-from paddle.inference import create_predictor
-from infer import Detector
+from .keypoint_preprocess import expand_crop
+from .keypoint_postprocess import HRNetPostProcess
+from .visualize import visualize_pose
+from .infer import Detector
 
 # Global dictionary
 KEYPOINT_SUPPORT_MODELS = {'HigherHRNet': 'keypoint_bottomup', 'HRNet': 'keypoint_topdown'}

modules/image/keypoint_detection/pp-tinypose/keypoint_postprocess.py

@@ -12,16 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from collections import abc
-from collections import defaultdict
 
 import cv2
 import numpy as np
-import paddle
-import paddle.nn as nn
-from keypoint_preprocess import get_affine_mat_kernel
-from keypoint_preprocess import get_affine_transform
-from scipy.optimize import linear_sum_assignment
+
+from .keypoint_preprocess import get_affine_transform
 
 
 class HRNetPostProcess(object):

modules/image/keypoint_detection/pp-tinypose/module.py

@@ -13,21 +13,19 @@
 # limitations under the License.
 import argparse
 import json
-import math
 import os
 import time
 from typing import Union
 
 import cv2
 import numpy as np
-import paddle
-import yaml
-from det_keypoint_unite_infer import predict_with_given_det
-from infer import Detector
-from keypoint_infer import KeyPointDetector
-from preprocess import base64_to_cv2
-from preprocess import decode_image
-from visualize import visualize_pose
+
+from .det_keypoint_unite_infer import predict_with_given_det
+from .infer import Detector
+from .keypoint_infer import KeyPointDetector
+from .preprocess import base64_to_cv2
+from .preprocess import decode_image
+from .visualize import visualize_pose
 
 from paddlehub.module.module import moduleinfo
 from paddlehub.module.module import runnable
@@ -40,7 +38,7 @@ from paddlehub.module.module import serving
     author="paddlepaddle",
     author_email="",
     summary="PP-TinyPose is a real-time keypoint detection model optimized by PaddleDetecion for mobile devices.",
-    version="1.0.0")
+    version="1.1.0")
 class PP_TinyPose:
     """
     PP-TinyPose Model.
@@ -74,12 +72,14 @@ class PP_TinyPose:
         results = self.detector.filter_box(results, 0.5)
         if results['boxes_num'] > 0:
             keypoint_res = predict_with_given_det(image, results, self.topdown_keypoint_detector, 1, False)
-            save_name = img if isinstance(img, str) else (str(time.time()) + '.png')
+            save_name = img if isinstance(img, str) else os.path.join(save_path, (str(time.time()) + '.png'))
             store_res.append(
                 [save_name, keypoint_res['bbox'], [keypoint_res['keypoint'][0], keypoint_res['keypoint'][1]]])
             if not os.path.exists(save_path):
                 os.makedirs(save_path)
             if self.visualization:
+                if isinstance(img, np.ndarray):
+                    cv2.imwrite(save_name, img)
                 visualize_pose(save_name, keypoint_res, visual_thresh=0.5, save_dir=save_path)
         return store_res
 

modules/image/keypoint_detection/pp-tinypose/preprocess.py

@@ -15,7 +15,8 @@ import base64
 
 import cv2
 import numpy as np
-from keypoint_preprocess import get_affine_transform
+
+from .keypoint_preprocess import get_affine_transform
 
 
 def decode_image(im_file, im_info):

modules/image/keypoint_detection/pp-tinypose/test.py

+import os
+import shutil
+import unittest
+
+import cv2
+import requests
+import paddlehub as hub
+
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+class TestHubModule(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        img_url = 'https://ai-studio-static-online.cdn.bcebos.com/7799a8ccc5f6471b9d56fb6eff94f82a08b70ca2c7594d3f99877e366c0a2619'
+        if not os.path.exists('tests'):
+            os.makedirs('tests')
+        response = requests.get(img_url)
+        assert response.status_code == 200, 'Network Error.'
+        with open('tests/test.jpg', 'wb') as f:
+            f.write(response.content)
+        cls.module = hub.Module(name="pp-tinypose")
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        shutil.rmtree('tests')
+        shutil.rmtree('pp_tinypose_output')
+
+    def test_predict1(self):
+        results = self.module.predict(
+            img='tests/test.jpg',
+            visualization=False
+        )
+        kps = results[0][1]
+        self.assertIsInstance(kps, list)
+
+    def test_predict2(self):
+        results = self.module.predict(
+            img=cv2.imread('tests/test.jpg'),
+            visualization=False
+        )
+        kps = results[0][1]
+        self.assertIsInstance(kps, list)
+
+    def test_predict3(self):
+        results = self.module.predict(
+            img=cv2.imread('tests/test.jpg'),
+            visualization=True
+        )
+        kps = results[0][1]
+        self.assertIsInstance(kps, list)
+
+    def test_predict4(self):
+        self.assertRaises(
+            FileNotFoundError,
+            self.module.predict,
+            img='no.jpg'
+        )
+
+    def test_predict5(self):
+        self.assertRaises(
+            cv2.error,
+            self.module.predict,
+            img=['test.jpg']
+        )
+
+if __name__ == "__main__":
+    unittest.main() 
\ No newline at end of file

modules/image/keypoint_detection/pp-tinypose/visualize.py

@@ -14,6 +14,7 @@
 from __future__ import division
 
 import os
+import math
 
 import cv2
 import numpy as np
@@ -22,7 +23,6 @@ from PIL import ImageDraw
 from PIL import ImageFile
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
-import math
 
 
 def visualize_box(im, results, labels, threshold=0.5):