pipeline.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  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, Optional, Union
  15. from ...common.batch_sampler import ImageBatchSampler
  16. from ...common.reader import ReadImage
  17. from ...utils.hpi import HPIConfig
  18. from ...utils.pp_option import PaddlePredictorOption
  19. from ..base import BasePipeline
  20. from ..components import CropByBoxes, FaissBuilder, FaissIndexer
  21. from .result import ShiTuResult
  22. class ShiTuV2Pipeline(BasePipeline):
  23. """ShiTuV2 Pipeline"""
  24. entities = "PP-ShiTuV2"
  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._topk, self._rec_threshold, self._hamming_radius, self._det_threshold = (
  37. config.get("rec_topk", 5),
  38. config.get("rec_threshold", 0.5),
  39. config.get("hamming_radius", None),
  40. config.get("det_threshold", 0.5),
  41. )
  42. index = config.get("index", None)
  43. self.img_reader = ReadImage(format="BGR")
  44. self.det_model = self.create_model(config["SubModules"]["Detection"])
  45. self.rec_model = self.create_model(config["SubModules"]["Recognition"])
  46. self.crop_by_boxes = CropByBoxes()
  47. self.indexer = FaissIndexer(index=index) if index else None
  48. self.batch_sampler = ImageBatchSampler(
  49. batch_size=self.det_model.batch_sampler.batch_size
  50. )
  51. def predict(self, input, index=None, **kwargs):
  52. indexer = FaissIndexer(index) if index is not None else self.indexer
  53. assert indexer
  54. kwargs = {k: v for k, v in kwargs.items() if v is not None}
  55. topk = kwargs.get("rec_topk", self._topk)
  56. rec_threshold = kwargs.get("rec_threshold", self._rec_threshold)
  57. hamming_radius = kwargs.get("hamming_radius", self._hamming_radius)
  58. det_threshold = kwargs.get("det_threshold", self._det_threshold)
  59. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  60. raw_imgs = self.img_reader(batch_data.instances)
  61. all_det_res = list(self.det_model(raw_imgs, threshold=det_threshold))
  62. for input_data, raw_img, det_res in zip(
  63. batch_data.instances, raw_imgs, all_det_res
  64. ):
  65. rec_res = self.get_rec_result(
  66. raw_img, det_res, indexer, rec_threshold, hamming_radius, topk
  67. )
  68. yield self.get_final_result(input_data, raw_img, det_res, rec_res)
  69. def get_rec_result(
  70. self, raw_img, det_res, indexer, rec_threshold, hamming_radius, topk
  71. ):
  72. if len(det_res["boxes"]) == 0:
  73. w, h = raw_img.shape[:2]
  74. det_res["boxes"].append(
  75. {
  76. "cls_id": 0,
  77. "label": "full_img",
  78. "score": 0,
  79. "coordinate": [0, 0, h, w],
  80. }
  81. )
  82. subs_of_img = list(self.crop_by_boxes(raw_img, det_res["boxes"]))
  83. img_list = [img["img"] for img in subs_of_img]
  84. all_rec_res = list(self.rec_model(img_list))
  85. all_rec_res = indexer(
  86. [rec_res["feature"] for rec_res in all_rec_res],
  87. score_thres=rec_threshold,
  88. hamming_radius=hamming_radius,
  89. topk=topk,
  90. )
  91. output = {"label": [], "score": []}
  92. for res in all_rec_res:
  93. output["label"].append(res["label"])
  94. output["score"].append(res["score"])
  95. return output
  96. def get_final_result(self, input_data, raw_img, det_res, rec_res):
  97. single_img_res = {"input_path": input_data, "input_img": raw_img, "boxes": []}
  98. for i, obj in enumerate(det_res["boxes"]):
  99. rec_scores = rec_res["score"][i]
  100. rec_scores = rec_scores if rec_scores is not None else [None]
  101. labels = rec_res["label"][i]
  102. labels = labels if labels is not None else [None]
  103. single_img_res["boxes"].append(
  104. {
  105. "labels": labels,
  106. "rec_scores": rec_scores,
  107. "det_score": obj["score"],
  108. "coordinate": obj["coordinate"],
  109. }
  110. )
  111. return ShiTuResult(single_img_res)
  112. def build_index(
  113. self,
  114. gallery_imgs,
  115. gallery_label,
  116. metric_type="IP",
  117. index_type="HNSW32",
  118. **kwargs
  119. ):
  120. return FaissBuilder.build(
  121. gallery_imgs,
  122. gallery_label,
  123. self.rec_model.predict,
  124. metric_type=metric_type,
  125. index_type=index_type,
  126. )
  127. def remove_index(self, remove_ids, index):
  128. return FaissBuilder.remove(remove_ids, index)
  129. def append_index(
  130. self,
  131. gallery_imgs,
  132. gallery_label,
  133. index,
  134. ):
  135. return FaissBuilder.append(
  136. gallery_imgs,
  137. gallery_label,
  138. self.rec_model.predict,
  139. index,
  140. )