Browse Source

Fix pipeline param bugs (#2921)

Zhang Zelun 10 months ago
parent
commit
ab4bddd451

+ 0 - 1
paddlex/inference/models_new/open_vocabulary_segmentation/results/sam_result.py

@@ -143,6 +143,5 @@ class SAMSegResult(BaseCVResult):
 
     def _to_json(self, *args, **kwargs):
         data = copy.deepcopy(self)
-        data["masks"] = "..."
         data.pop("input_img")
         return JsonMixin._to_json(data, *args, **kwargs)

+ 7 - 3
paddlex/inference/pipelines_new/instance_segmentation/pipeline.py

@@ -53,15 +53,19 @@ class InstanceSegmentationPipeline(BasePipeline):
         self.threshold = instance_segmentation_model_config["threshold"]
 
     def predict(
-        self, input: Union[str, List[str], np.ndarray, List[np.ndarray]], **kwargs
+        self,
+        input: str | list[str] | np.ndarray | list[np.ndarray],
+        threshold: float | None = None,
+        **kwargs
     ) -> InstanceSegResult:
         """Predicts instance segmentation results for the given input.
 
         Args:
-            input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
+            input (str | list[str] | np.ndarray | list[np.ndarray]): The input image(s) or path(s) to the images.
+            threshold (float | None): The threshold value to filter out low-confidence predictions. Default is None.
             **kwargs: Additional keyword arguments that can be passed to the function.
 
         Returns:
             InstanceSegResult: The predicted instance segmentation results.
         """
-        yield from self.instance_segmentation_model(input, threshold=self.threshold)
+        yield from self.instance_segmentation_model(input, threshold=threshold)

+ 3 - 1
paddlex/inference/pipelines_new/open_vocabulary_detection/pipeline.py

@@ -57,6 +57,7 @@ class OpenVocabularyDetectionPipeline(BasePipeline):
         self,
         input: Union[str, List[str], np.ndarray, List[np.ndarray]],
         prompt: str,
+        thresholds: dict[str, float] | None = None,
         **kwargs
     ) -> DetResult:
         """Predicts open vocabulary detection results for the given input.
@@ -64,11 +65,12 @@ class OpenVocabularyDetectionPipeline(BasePipeline):
         Args:
             input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
             prompt (str): The text prompt used to describe the objects.
+            thresholds (dict | None): Threshold values for different models. If provided, these will override any default threshold values set during initialization. Default is None.
             **kwargs: Additional keyword arguments that can be passed to the function.
 
         Returns:
             DetResult: The predicted open vocabulary detection results.
         """
         yield from self.open_vocabulary_detection_model(
-            input, prompt=prompt, thresholds=self.thresholds
+            input, prompt=prompt, thresholds=thresholds
         )

+ 2 - 2
paddlex/inference/pipelines_new/open_vocabulary_segmentation/pipeline.py

@@ -70,8 +70,8 @@ class OpenVocabularySegmentationPipeline(BasePipeline):
         """Predicts image segmentation results for the given input.
 
         Args:
-            input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
-            prompt (list[list[int]] | np.ndarray): The prompt for the input image(s).
+            input (str | list[str] | np.ndarray | list[np.ndarray]): The input image(s) or path(s) to the images.
+            prompt (list[list[float]] | np.ndarray): The prompt for the input image(s).
             prompt_type (str): The type of prompt, either 'box' or 'point'. Default is 'box'.
             **kwargs: Additional keyword arguments that can be passed to the function.
 

+ 11 - 3
paddlex/inference/pipelines_new/rotated_object_detection/pipeline.py

@@ -53,15 +53,23 @@ class RotatedObjectDetectionPipeline(BasePipeline):
         self.threshold = rotated_object_detection_model_config["threshold"]
 
     def predict(
-        self, input: Union[str, List[str], np.ndarray, List[np.ndarray]], **kwargs
+        self,
+        input: str | list[str] | np.ndarray | list[np.ndarray],
+        threshold: None | dict[int, float] | float = None,
+        **kwargs
     ) -> DetResult:
         """Predicts rotated object detection results for the given input.
 
         Args:
-            input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
+            input (str | list[str] | np.ndarray | list[np.ndarray]): The input image(s) or path(s) to the images.
+            threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
+                If None, it will use the default threshold specified during initialization.
+                If a dictionary is provided, it should have integer keys corresponding to the class IDs and float values
+                representing the respective thresholds for each class.
+                If a single float value is provided, it will be used as the threshold for all classes.
             **kwargs: Additional keyword arguments that can be passed to the function.
 
         Returns:
             DetResult: The predicted rotated object detection results.
         """
-        yield from self.rotated_object_detection_model(input, threshold=self.threshold)
+        yield from self.rotated_object_detection_model(input, threshold=threshold)

+ 12 - 4
paddlex/inference/pipelines_new/semantic_segmentation/pipeline.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Union, List
+from typing import Any, Dict, Optional, Literal
 import numpy as np
 from ...utils.pp_option import PaddlePredictorOption
 from ..base import BasePipeline
@@ -53,15 +53,23 @@ class SemanticSegmentationPipeline(BasePipeline):
         self.target_size = semantic_segmentation_model_config["target_size"]
 
     def predict(
-        self, input: Union[str, List[str], np.ndarray, List[np.ndarray]], **kwargs
+        self,
+        input: str | list[str] | np.ndarray | list[np.ndarray],
+        target_size: Literal[-1] | None | int | tuple[int] = None,
+        **kwargs
     ) -> SegResult:
         """Predicts semantic segmentation results for the given input.
 
         Args:
-            input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
+            input (str | list[str] | np.ndarray | list[np.ndarray]): The input image(s) or path(s) to the images.
+            target_size (Literal[-1] | None | int | tuple[int]): The Image size model used to do prediction. Default is None.
+                If it's set to -1, the original image size will be used.
+                If it's set to None, the previous level's setting will be used.
+                If it's set to an integer value, the image will be rescaled to the size of (value, value).
+                If it's set to a tuple of two integers, the image will be rescaled to the size of (height, width).
             **kwargs: Additional keyword arguments that can be passed to the function.
 
         Returns:
             SegResult: The predicted segmentation results.
         """
-        yield from self.semantic_segmentation_model(input, target_size=self.target_size)
+        yield from self.semantic_segmentation_model(input, target_size=target_size)

+ 11 - 3
paddlex/inference/pipelines_new/small_object_detection/pipeline.py

@@ -53,15 +53,23 @@ class SmallObjectDetectionPipeline(BasePipeline):
         self.threshold = small_object_detection_model_config["threshold"]
 
     def predict(
-        self, input: Union[str, List[str], np.ndarray, List[np.ndarray]], **kwargs
+        self,
+        input: str | list[str] | np.ndarray | list[np.ndarray],
+        threshold: None | dict[int, float] | float = None,
+        **kwargs
     ) -> DetResult:
         """Predicts small object detection results for the given input.
 
         Args:
-            input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
+            input (str | list[str] | np.ndarray | list[np.ndarray]): The input image(s) or path(s) to the images.
+            threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
+                If None, it will use the default threshold specified during initialization.
+                If a dictionary is provided, it should have integer keys corresponding to the class IDs and float values
+                representing the respective thresholds for each class.
+                If a single float value is provided, it will be used as the threshold for all classes.
             **kwargs: Additional keyword arguments that can be passed to the function.
 
         Returns:
             DetResult: The predicted small object detection results.
         """
-        yield from self.small_object_detection_model(input, threshold=self.threshold)
+        yield from self.small_object_detection_model(input, threshold=threshold)