pp_shitu_v2.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178
  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 pathlib import Path
  15. import numpy as np
  16. from ..utils.io import ImageReader
  17. from ..components import CropByBoxes, FaissIndexer
  18. from ..components.retrieval.faiss import FaissBuilder
  19. from ..results import ShiTuResult
  20. from .base import BasePipeline
  21. class ShiTuV2Pipeline(BasePipeline):
  22. """ShiTuV2 Pipeline"""
  23. entities = "PP-ShiTuV2"
  24. def __init__(
  25. self,
  26. det_model,
  27. rec_model,
  28. det_batch_size=1,
  29. rec_batch_size=1,
  30. index_dir=None,
  31. metric_type="IP",
  32. score_thres=None,
  33. hamming_radius=None,
  34. return_k=5,
  35. device=None,
  36. predictor_kwargs=None,
  37. ):
  38. super().__init__(device, predictor_kwargs)
  39. self._build_predictor(det_model, rec_model)
  40. self.set_predictor(det_batch_size, rec_batch_size, device)
  41. self._metric_type, self._return_k, self._score_thres, self._hamming_radius = (
  42. metric_type,
  43. return_k,
  44. score_thres,
  45. hamming_radius,
  46. )
  47. self._indexer = self._build_indexer(index_dir) if index_dir else None
  48. def _build_indexer(self, index_dir):
  49. return FaissIndexer(
  50. index_dir,
  51. self._metric_type,
  52. self._return_k,
  53. self._score_thres,
  54. self._hamming_radius,
  55. )
  56. def _build_predictor(self, det_model, rec_model):
  57. self.det_model = self._create(model=det_model)
  58. self.rec_model = self._create(model=rec_model)
  59. self._crop_by_boxes = CropByBoxes()
  60. self._img_reader = ImageReader(backend="opencv")
  61. def set_predictor(self, det_batch_size=None, rec_batch_size=None, device=None):
  62. if det_batch_size:
  63. self.det_model.set_predictor(batch_size=det_batch_size)
  64. if rec_batch_size:
  65. self.rec_model.set_predictor(batch_size=rec_batch_size)
  66. if device:
  67. self.det_model.set_predictor(device=device)
  68. self.rec_model.set_predictor(device=device)
  69. def predict(self, input, index_dir=None, **kwargs):
  70. indexer = self._build_indexer(index_dir) if index_dir else self._indexer
  71. assert indexer
  72. self.set_predictor(**kwargs)
  73. for det_res in self.det_model(input):
  74. rec_res = self.get_rec_result(det_res, indexer)
  75. yield self.get_final_result(det_res, rec_res)
  76. def get_rec_result(self, det_res, indexer):
  77. full_img = self._img_reader.read(det_res["input_path"])
  78. w, h = full_img.shape[:2]
  79. det_res["boxes"].append(
  80. {"cls_id": 0, "label": "full_img", "score": 0, "coordinate": [0, 0, h, w]}
  81. )
  82. subs_of_img = list(self._crop_by_boxes(det_res))
  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 = next(indexer(all_rec_res))
  86. output = {"label": [], "score": []}
  87. for res in all_rec_res:
  88. output["label"].append(res["label"])
  89. output["score"].append(res["score"])
  90. return output
  91. def get_final_result(self, det_res, rec_res):
  92. single_img_res = {"input_path": det_res["input_path"], "boxes": []}
  93. for i, obj in enumerate(det_res["boxes"]):
  94. rec_scores = rec_res["score"][i]
  95. labels = rec_res["label"][i]
  96. single_img_res["boxes"].append(
  97. {
  98. "labels": labels,
  99. "rec_scores": rec_scores,
  100. "det_score": obj["score"],
  101. "coordinate": obj["coordinate"],
  102. }
  103. )
  104. return ShiTuResult(single_img_res)
  105. def _build_index(
  106. self,
  107. data_root,
  108. index_dir,
  109. mode="new",
  110. metric_type="IP",
  111. index_type="HNSW32",
  112. **kwargs,
  113. ):
  114. self.set_predictor(**kwargs)
  115. self._metric_type = metric_type if metric_type else self._metric_type
  116. builder = FaissBuilder(
  117. self.rec_model.predict,
  118. mode=mode,
  119. metric_type=self._metric_type,
  120. index_type=index_type,
  121. )
  122. if mode == "new":
  123. builder.build(Path(data_root) / "gallery.txt", data_root, index_dir)
  124. elif mode == "remove":
  125. builder.remove(Path(data_root) / "gallery.txt", data_root, index_dir)
  126. elif mode == "append":
  127. builder.append(Path(data_root) / "gallery.txt", data_root, index_dir)
  128. else:
  129. raise Exception("`mode` only support `new`, `remove` and `append`.")
  130. def build_index(
  131. self, data_root, index_dir, metric_type="IP", index_type="HNSW32", **kwargs
  132. ):
  133. self._build_index(
  134. data_root=data_root,
  135. index_dir=index_dir,
  136. mode="new",
  137. metric_type=metric_type,
  138. index_type=index_type,
  139. **kwargs,
  140. )
  141. def remove_index(
  142. self, data_root, index_dir, metric_type="IP", index_type="HNSW32", **kwargs
  143. ):
  144. self._build_index(
  145. data_root=data_root,
  146. index_dir=index_dir,
  147. mode="remove",
  148. metric_type=metric_type,
  149. index_type=index_type,
  150. **kwargs,
  151. )
  152. def append_index(
  153. self, data_root, index_dir, metric_type="IP", index_type="HNSW32", **kwargs
  154. ):
  155. self._build_index(
  156. data_root=data_root,
  157. index_dir=index_dir,
  158. mode="append",
  159. metric_type=metric_type,
  160. index_type=index_type,
  161. **kwargs,
  162. )