Kaynağa Gözat

fix cls pipline

zhangyubo0722 10 ay önce
ebeveyn
işleme
e9cfba7cf0

+ 2 - 1
paddlex/configs/pipelines/image_classification.yaml

@@ -6,4 +6,5 @@ SubModules:
     module_name: image_classification
     model_name: PP-LCNet_x0_5
     model_dir: null
-    batch_size: 4    
+    batch_size: 4
+    topk: 5 

+ 2 - 8
paddlex/inference/pipelines_new/image_classification/pipeline.py

@@ -54,9 +54,7 @@ class ImageClassificationPipeline(BasePipeline):
         self.image_classification_model = self.create_model(
             image_classification_model_config
         )
-        batch_size = image_classification_model_config["batch_size"]
-        self.batch_sampler = ImageBatchSampler(batch_size=batch_size)
-        self.img_reader = ReadImage(format="BGR")
+        self.topk = image_classification_model_config["topk"]
 
     def predict(
         self, input: str | list[str] | np.ndarray | list[np.ndarray], **kwargs
@@ -70,8 +68,4 @@ class ImageClassificationPipeline(BasePipeline):
         Returns:
             TopkResult: The predicted top k results.
         """
-
-        for img_id, batch_data in enumerate(self.batch_sampler(input)):
-            batch_imgs = self.img_reader(batch_data)
-            for topk_single_result in self.image_classification_model(batch_imgs):
-                yield topk_single_result
+        yield from self.image_classification_model(input, topk=self.topk)