pp_shitu_v2.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  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. import pickle
  15. from pathlib import Path
  16. import numpy as np
  17. from ..utils.io import ImageReader
  18. from ..components import CropByBoxes, FaissIndexer
  19. from ..components.retrieval.faiss import FaissBuilder
  20. from ..results import ShiTuResult
  21. from .base import BasePipeline
  22. class ShiTuV2Pipeline(BasePipeline):
  23. """ShiTuV2 Pipeline"""
  24. entities = "PP-ShiTuV2"
  25. def __init__(
  26. self,
  27. det_model,
  28. rec_model,
  29. det_batch_size=1,
  30. rec_batch_size=1,
  31. index=None,
  32. metric_type="IP",
  33. score_thres=None,
  34. hamming_radius=None,
  35. return_k=5,
  36. device=None,
  37. predictor_kwargs=None,
  38. ):
  39. super().__init__(device, predictor_kwargs)
  40. self._build_predictor(det_model, rec_model)
  41. self.set_predictor(det_batch_size, rec_batch_size, device)
  42. self._metric_type, self._return_k, self._score_thres, self._hamming_radius = (
  43. metric_type,
  44. return_k,
  45. score_thres,
  46. hamming_radius,
  47. )
  48. self._indexer = self._build_indexer(index=index) if index else None
  49. def _build_indexer(self, index):
  50. return FaissIndexer(
  51. index=index,
  52. metric_type=self._metric_type,
  53. return_k=self._return_k,
  54. score_thres=self._score_thres,
  55. hamming_radius=self._hamming_radius,
  56. )
  57. def _build_predictor(self, det_model, rec_model):
  58. self.det_model = self._create(model=det_model)
  59. self.rec_model = self._create(model=rec_model)
  60. self._crop_by_boxes = CropByBoxes()
  61. self._img_reader = ImageReader(backend="opencv")
  62. def set_predictor(self, det_batch_size=None, rec_batch_size=None, device=None):
  63. if det_batch_size:
  64. self.det_model.set_predictor(batch_size=det_batch_size)
  65. if rec_batch_size:
  66. self.rec_model.set_predictor(batch_size=rec_batch_size)
  67. if device:
  68. self.det_model.set_predictor(device=device)
  69. self.rec_model.set_predictor(device=device)
  70. def predict(self, input, index=None, **kwargs):
  71. indexer = self._build_indexer(index) if index is not None else self._indexer
  72. assert indexer
  73. self.set_predictor(**kwargs)
  74. for det_res in self.det_model(input):
  75. rec_res = self.get_rec_result(det_res, indexer)
  76. yield self.get_final_result(det_res, rec_res)
  77. def get_rec_result(self, det_res, indexer):
  78. if len(det_res["boxes"]) == 0:
  79. full_img = self._img_reader.read(det_res["input_path"])
  80. w, h = full_img.shape[:2]
  81. det_res["boxes"].append(
  82. {
  83. "cls_id": 0,
  84. "label": "full_img",
  85. "score": 0,
  86. "coordinate": [0, 0, h, w],
  87. }
  88. )
  89. subs_of_img = list(self._crop_by_boxes(det_res))
  90. img_list = [img["img"] for img in subs_of_img]
  91. all_rec_res = list(self.rec_model(img_list))
  92. all_rec_res = next(indexer(all_rec_res))
  93. output = {"label": [], "score": []}
  94. for res in all_rec_res:
  95. output["label"].append(res["label"])
  96. output["score"].append(res["score"])
  97. return output
  98. def get_final_result(self, det_res, rec_res):
  99. single_img_res = {"input_path": det_res["input_path"], "boxes": []}
  100. for i, obj in enumerate(det_res["boxes"]):
  101. rec_scores = rec_res["score"][i]
  102. labels = rec_res["label"][i]
  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. **kwargs
  127. )
  128. def remove_index(
  129. self, gallery_imgs, gallery_label, index, index_type="HNSW32", **kwargs
  130. ):
  131. return FaissBuilder.remove(
  132. gallery_imgs, gallery_label, index, index_type=index_type, **kwargs
  133. )
  134. def append_index(
  135. self,
  136. gallery_imgs,
  137. gallery_label,
  138. index,
  139. id_map=None,
  140. metric_type="IP",
  141. index_type="HNSW32",
  142. **kwargs
  143. ):
  144. return FaissBuilder.append(
  145. gallery_imgs,
  146. gallery_label,
  147. self.rec_model.predict,
  148. index,
  149. metric_type=metric_type,
  150. **kwargs
  151. )