pipeline.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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, List, Optional, Union
  15. import numpy as np
  16. from ...common.batch_sampler import ImageBatchSampler
  17. from ...common.reader import ReadImage
  18. from ...utils.hpi import HPIConfig
  19. from ...utils.pp_option import PaddlePredictorOption
  20. from ..base import BasePipeline
  21. from ..components import CropByBoxes
  22. from .result import AttributeRecResult
  23. class AttributeRecPipeline(BasePipeline):
  24. """Attribute Rec Pipeline"""
  25. def __init__(
  26. self,
  27. config: Dict,
  28. device: str = None,
  29. pp_option: PaddlePredictorOption = None,
  30. use_hpip: bool = False,
  31. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  32. ):
  33. super().__init__(
  34. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  35. )
  36. self.det_model = self.create_model(config["SubModules"]["Detection"])
  37. self.cls_model = self.create_model(config["SubModules"]["Classification"])
  38. self._crop_by_boxes = CropByBoxes()
  39. self._img_reader = ReadImage(format="BGR")
  40. self.det_threshold = config["SubModules"]["Detection"].get("threshold", 0.5)
  41. self.cls_threshold = config["SubModules"]["Classification"].get(
  42. "threshold", 0.7
  43. )
  44. self.batch_sampler = ImageBatchSampler(
  45. batch_size=config["SubModules"]["Detection"]["batch_size"]
  46. )
  47. self.img_reader = ReadImage(format="BGR")
  48. def predict(
  49. self,
  50. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  51. det_threshold: float = None,
  52. cls_threshold: Union[float, dict, list, None] = None,
  53. **kwargs
  54. ):
  55. det_threshold = self.det_threshold if det_threshold is None else det_threshold
  56. cls_threshold = self.cls_threshold if cls_threshold is None else cls_threshold
  57. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  58. raw_imgs = self.img_reader(batch_data.instances)
  59. all_det_res = list(self.det_model(raw_imgs, threshold=det_threshold))
  60. for input_path, input_data, raw_img, det_res in zip(
  61. batch_data.input_paths, batch_data.instances, raw_imgs, all_det_res
  62. ):
  63. cls_res = self.get_cls_result(raw_img, det_res, cls_threshold)
  64. yield self.get_final_result(input_path, raw_img, det_res, cls_res)
  65. def get_cls_result(self, raw_img, det_res, cls_threshold):
  66. subs_of_img = list(self._crop_by_boxes(raw_img, det_res["boxes"]))
  67. img_list = [img["img"] for img in subs_of_img]
  68. all_cls_res = list(self.cls_model(img_list, threshold=cls_threshold))
  69. output = {"label": [], "score": []}
  70. for res in all_cls_res:
  71. output["label"].append(res["label_names"])
  72. output["score"].append(res["scores"])
  73. return output
  74. def get_final_result(self, input_path, raw_img, det_res, rec_res):
  75. single_img_res = {"input_path": input_path, "input_img": raw_img, "boxes": []}
  76. for i, obj in enumerate(det_res["boxes"]):
  77. cls_scores = rec_res["score"][i]
  78. labels = rec_res["label"][i]
  79. single_img_res["boxes"].append(
  80. {
  81. "labels": labels,
  82. "cls_scores": cls_scores,
  83. "det_score": obj["score"],
  84. "coordinate": obj["coordinate"],
  85. }
  86. )
  87. return AttributeRecResult(single_img_res)
  88. class PedestrianAttributeRecPipeline(AttributeRecPipeline):
  89. entities = "pedestrian_attribute_recognition"
  90. class VehicleAttributeRecPipeline(AttributeRecPipeline):
  91. entities = "vehicle_attribute_recognition"