소스 검색

bugfix: kpt model support to pass dict as input

gaotingquan 6 달 전
부모
커밋
a589bbbe84
2개의 변경된 파일6개의 추가작업 그리고 1개의 파일을 삭제
  1. 3 1
      paddlex/inference/common/batch_sampler/image_batch_sampler.py
  2. 3 0
      paddlex/inference/models/keypoint_detection/predictor.py

+ 3 - 1
paddlex/inference/common/batch_sampler/image_batch_sampler.py

@@ -56,7 +56,9 @@ class ImageBatchSampler(BaseBatchSampler):
 
     def _get_files_list(self, fp):
         if fp is None or not os.path.exists(fp):
-            raise Exception(f"Not found any image files or pdf files in path: {fp}")
+            raise Exception(f"Not found any files in path: {fp}")
+        if os.path.isfile(fp):
+            return [fp]
 
         file_list = []
         if os.path.isdir(fp):

+ 3 - 0
paddlex/inference/models/keypoint_detection/predictor.py

@@ -26,6 +26,9 @@ from .result import KptResult
 
 
 class KptBatchSampler(ImageBatchSampler):
+    # don't support to pass pdf file as input
+    PDF_SUFFIX = []
+
     def sample(self, inputs):
         if not isinstance(inputs, list):
             inputs = [inputs]