pipeline.py 5.3 KB

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