face_recognition.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  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 .pp_shitu_v2 import ShiTuV2Pipeline
  15. from ..results import FaceRecResult
  16. class FaceRecPipeline(ShiTuV2Pipeline):
  17. """Face Recognition Pipeline"""
  18. entities = "face_recognition"
  19. def get_rec_result(self, det_res, indexer):
  20. if len(det_res["boxes"]) == 0:
  21. return {"label": [], "score": []}
  22. subs_of_img = list(self._crop_by_boxes(det_res))
  23. img_list = [img["img"] for img in subs_of_img]
  24. all_rec_res = list(self.rec_model(img_list))
  25. all_rec_res = next(indexer(all_rec_res))
  26. output = {"label": [], "score": []}
  27. for res in all_rec_res:
  28. output["label"].append(res["label"])
  29. output["score"].append(res["score"])
  30. return output
  31. def get_final_result(self, det_res, rec_res):
  32. single_img_res = {"input_path": det_res["input_path"], "boxes": []}
  33. for i, obj in enumerate(det_res["boxes"]):
  34. rec_scores = rec_res["score"][i]
  35. labels = rec_res["label"][i]
  36. single_img_res["boxes"].append(
  37. {
  38. "labels": labels,
  39. "rec_scores": rec_scores,
  40. "det_score": obj["score"],
  41. "coordinate": obj["coordinate"],
  42. }
  43. )
  44. return FaceRecResult(single_img_res)