face_recognition.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  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. subs_of_img = list(self._crop_by_boxes(det_res))
  21. img_list = [img["img"] for img in subs_of_img]
  22. all_rec_res = list(self.rec_model(img_list))
  23. all_rec_res = next(indexer(all_rec_res))
  24. output = {"label": [], "score": []}
  25. for res in all_rec_res:
  26. output["label"].append(res["label"])
  27. output["score"].append(res["score"])
  28. return output
  29. def get_final_result(self, det_res, rec_res):
  30. single_img_res = {"input_path": det_res["input_path"], "boxes": []}
  31. for i, obj in enumerate(det_res["boxes"]):
  32. rec_scores = rec_res["score"][i]
  33. labels = rec_res["label"][i]
  34. single_img_res["boxes"].append(
  35. {
  36. "labels": labels,
  37. "rec_scores": rec_scores,
  38. "det_score": obj["score"],
  39. "coordinate": obj["coordinate"],
  40. }
  41. )
  42. return FaceRecResult(single_img_res)