pipeline.py 2.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Any, Dict, Optional, Union, List
  15. import numpy as np
  16. from ...utils.pp_option import PaddlePredictorOption
  17. from ..base import BasePipeline
  18. # [TODO] 待更新models_new到models
  19. from ...models_new.video_detection.result import DetVideoResult
  20. class VideoDetectionPipeline(BasePipeline):
  21. """Video detection Pipeline"""
  22. entities = "video_detection"
  23. def __init__(
  24. self,
  25. config: Dict,
  26. device: str = None,
  27. pp_option: PaddlePredictorOption = None,
  28. use_hpip: bool = False,
  29. ) -> None:
  30. """
  31. Initializes the class with given configurations and options.
  32. Args:
  33. config (Dict): Configuration dictionary containing model and other parameters.
  34. device (str): The device to run the prediction on. Default is None.
  35. pp_option (PaddlePredictorOption): Options for PaddlePaddle predictor. Default is None.
  36. use_hpip (bool): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
  37. """
  38. super().__init__(device=device, pp_option=pp_option, use_hpip=use_hpip)
  39. video_detection_model_config = config["SubModules"]["VideoDetection"]
  40. model_kwargs = {}
  41. if "nms_thresh" in video_detection_model_config:
  42. model_kwargs["nms_thresh"] = video_detection_model_config["nms_thresh"]
  43. if "score_thresh" in video_detection_model_config:
  44. model_kwargs["score_thresh"] = video_detection_model_config["score_thresh"]
  45. self.video_detection_model = self.create_model(video_detection_model_config, **model_kwargs)
  46. def predict(
  47. self,
  48. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  49. nms_thresh: float = 0.5,
  50. score_thresh: float = 0.4,
  51. **kwargs
  52. ) -> DetVideoResult:
  53. """Predicts video detection results for the given input.
  54. Args:
  55. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or path(s) to the images.
  56. **kwargs: Additional keyword arguments that can be passed to the function.
  57. Returns:
  58. DetVideoResult: The predicted video detection results.
  59. """
  60. yield from self.video_detection_model(
  61. input, nms_thresh=nms_thresh, score_thresh=score_thresh
  62. )