face_recognition.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312
  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 asyncio
  15. import faiss
  16. import pickle
  17. import uuid
  18. from typing import Dict, List, Optional
  19. from fastapi import FastAPI, HTTPException
  20. from pydantic import BaseModel, Field
  21. from typing_extensions import Annotated, TypeAlias
  22. from .....utils import logging
  23. from ....components.retrieval.faiss import IndexData
  24. from ...face_recognition import FaceRecPipeline
  25. from ..storage import create_storage
  26. from .. import utils as serving_utils
  27. from ..app import AppConfig, create_app
  28. from ..models import NoResultResponse, ResultResponse
  29. class ImageLabelPair(BaseModel):
  30. image: str
  31. label: str
  32. class BuildIndexRequest(BaseModel):
  33. imageLabelPairs: List[ImageLabelPair]
  34. class BuildIndexResult(BaseModel):
  35. indexKey: str
  36. idMap: Dict[int, str]
  37. class AddImagesToIndexRequest(BaseModel):
  38. imageLabelPairs: List[ImageLabelPair]
  39. indexKey: str
  40. class AddImagesToIndexResult(BaseModel):
  41. idMap: Dict[int, str]
  42. class RemoveImagesFromIndexRequest(BaseModel):
  43. ids: List[int]
  44. indexKey: str
  45. class RemoveImagesFromIndexResult(BaseModel):
  46. idMap: Dict[int, str]
  47. class InferRequest(BaseModel):
  48. image: str
  49. indexKey: Optional[str] = None
  50. BoundingBox: TypeAlias = Annotated[List[float], Field(min_length=4, max_length=4)]
  51. class RecResult(BaseModel):
  52. label: str
  53. score: float
  54. class Face(BaseModel):
  55. bbox: BoundingBox
  56. recResults: List[RecResult]
  57. score: float
  58. class InferResult(BaseModel):
  59. faces: List[Face]
  60. image: str
  61. def _serialize_index_data(index_data: IndexData) -> bytes:
  62. tup = (index_data.index_bytes, index_data.index_info)
  63. return pickle.dumps(tup)
  64. def _deserialize_index_data(index_data_bytes: bytes) -> IndexData:
  65. tup = pickle.loads(index_data_bytes)
  66. index = faiss.deserialize_index(tup[0])
  67. return IndexData(index, tup[1])
  68. def create_pipeline_app(pipeline: FaceRecPipeline, app_config: AppConfig) -> FastAPI:
  69. app, ctx = create_app(
  70. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  71. )
  72. if ctx.config.extra and "index_storage" in ctx.config.extra:
  73. ctx.extra["index_storage"] = create_storage(ctx.config.extra["index_storage"])
  74. else:
  75. ctx.extra["index_storage"] = create_storage({"type": "memory"})
  76. @app.post(
  77. "/face-recognition-index-build",
  78. operation_id="buildIndex",
  79. responses={422: {"model": NoResultResponse}},
  80. response_model_exclude_none=True,
  81. )
  82. async def _build_index(
  83. request: BuildIndexRequest,
  84. ) -> ResultResponse[BuildIndexResult]:
  85. pipeline = ctx.pipeline
  86. aiohttp_session = ctx.aiohttp_session
  87. try:
  88. images = [pair.image for pair in request.imageLabelPairs]
  89. file_bytes_list = await asyncio.gather(
  90. *(serving_utils.get_raw_bytes(img, aiohttp_session) for img in images)
  91. )
  92. images = [
  93. serving_utils.image_bytes_to_array(item) for item in file_bytes_list
  94. ]
  95. labels = [pair.label for pair in request.imageLabelPairs]
  96. # TODO: Support specifying `index_type` and `metric_type` in the
  97. # request
  98. index_data = await pipeline.call(
  99. pipeline.pipeline.build_index,
  100. images,
  101. labels,
  102. index_type="Flat",
  103. metric_type="IP",
  104. )
  105. index_storage = ctx.extra["index_storage"]
  106. index_key = str(uuid.uuid4())
  107. index_data_bytes = await serving_utils.call_async(
  108. _serialize_index_data, index_data
  109. )
  110. await serving_utils.call_async(
  111. index_storage.set, index_key, index_data_bytes
  112. )
  113. return ResultResponse[BuildIndexResult](
  114. logId=serving_utils.generate_log_id(),
  115. result=BuildIndexResult(indexKey=index_key, idMap=index_data.id_map),
  116. )
  117. except Exception:
  118. logging.exception("Unexpected exception")
  119. raise HTTPException(status_code=500, detail="Internal server error")
  120. @app.post(
  121. "/face-recognition-index-add",
  122. operation_id="buildIndex",
  123. responses={422: {"model": NoResultResponse}},
  124. response_model_exclude_none=True,
  125. )
  126. async def _add_images_to_index(
  127. request: AddImagesToIndexRequest,
  128. ) -> ResultResponse[AddImagesToIndexResult]:
  129. pipeline = ctx.pipeline
  130. aiohttp_session = ctx.aiohttp_session
  131. try:
  132. images = [pair.image for pair in request.imageLabelPairs]
  133. file_bytes_list = await asyncio.gather(
  134. *(serving_utils.get_raw_bytes(img, aiohttp_session) for img in images)
  135. )
  136. images = [
  137. serving_utils.image_bytes_to_array(item) for item in file_bytes_list
  138. ]
  139. labels = [pair.label for pair in request.imageLabelPairs]
  140. index_storage = ctx.extra["index_storage"]
  141. index_data_bytes = await serving_utils.call_async(
  142. index_storage.get, request.indexKey
  143. )
  144. index_data = await serving_utils.call_async(
  145. _deserialize_index_data, index_data_bytes
  146. )
  147. index_data = await pipeline.call(
  148. pipeline.pipeline.append_index, images, labels, index_data
  149. )
  150. index_data_bytes = await serving_utils.call_async(
  151. _serialize_index_data, index_data
  152. )
  153. await serving_utils.call_async(
  154. index_storage.set, request.indexKey, index_data_bytes
  155. )
  156. return ResultResponse[AddImagesToIndexResult](
  157. logId=serving_utils.generate_log_id(),
  158. result=AddImagesToIndexResult(idMap=index_data.id_map),
  159. )
  160. except Exception:
  161. logging.exception("Unexpected exception")
  162. raise HTTPException(status_code=500, detail="Internal server error")
  163. @app.post(
  164. "/face-recognition-index-remove",
  165. operation_id="buildIndex",
  166. responses={422: {"model": NoResultResponse}},
  167. response_model_exclude_none=True,
  168. )
  169. async def _remove_images_from_index(
  170. request: RemoveImagesFromIndexRequest,
  171. ) -> ResultResponse[RemoveImagesFromIndexResult]:
  172. pipeline = ctx.pipeline
  173. try:
  174. index_storage = ctx.extra["index_storage"]
  175. index_data_bytes = await serving_utils.call_async(
  176. index_storage.get, request.indexKey
  177. )
  178. index_data = await serving_utils.call_async(
  179. _deserialize_index_data, index_data_bytes
  180. )
  181. index_data = await pipeline.call(
  182. pipeline.pipeline.remove_index, request.ids, index_data
  183. )
  184. index_data_bytes = await serving_utils.call_async(
  185. _serialize_index_data, index_data
  186. )
  187. await serving_utils.call_async(
  188. index_storage.set, request.indexKey, index_data_bytes
  189. )
  190. return ResultResponse[RemoveImagesFromIndexResult](
  191. logId=serving_utils.generate_log_id(),
  192. result=RemoveImagesFromIndexResult(idMap=index_data.id_map),
  193. )
  194. except Exception:
  195. logging.exception("Unexpected exception")
  196. raise HTTPException(status_code=500, detail="Internal server error")
  197. @app.post(
  198. "/face-recognition-infer",
  199. operation_id="infer",
  200. responses={422: {"model": NoResultResponse}},
  201. response_model_exclude_none=True,
  202. )
  203. async def _infer(request: InferRequest) -> ResultResponse[InferResult]:
  204. pipeline = ctx.pipeline
  205. aiohttp_session = ctx.aiohttp_session
  206. try:
  207. image_bytes = await serving_utils.get_raw_bytes(
  208. request.image, aiohttp_session
  209. )
  210. image = serving_utils.image_bytes_to_array(image_bytes)
  211. if request.indexKey is not None:
  212. index_storage = ctx.extra["index_storage"]
  213. index_data_bytes = await serving_utils.call_async(
  214. index_storage.get, request.indexKey
  215. )
  216. index_data = await serving_utils.call_async(
  217. _deserialize_index_data, index_data_bytes
  218. )
  219. else:
  220. index_data = None
  221. result = list(
  222. await pipeline.call(pipeline.pipeline.predict, image, index=index_data)
  223. )[0]
  224. faces: List[Face] = []
  225. for face in result["boxes"]:
  226. rec_results: List[RecResult] = []
  227. if face["rec_scores"] is not None:
  228. for label, score in zip(face["labels"], face["rec_scores"]):
  229. rec_results.append(
  230. RecResult(
  231. label=label,
  232. score=score,
  233. )
  234. )
  235. faces.append(
  236. Face(
  237. bbox=face["coordinate"],
  238. recResults=rec_results,
  239. score=face["det_score"],
  240. )
  241. )
  242. output_image_base64 = serving_utils.base64_encode(
  243. serving_utils.image_to_bytes(result.img)
  244. )
  245. return ResultResponse[InferResult](
  246. logId=serving_utils.generate_log_id(),
  247. result=InferResult(faces=faces, image=output_image_base64),
  248. )
  249. except Exception:
  250. logging.exception("Unexpected exception")
  251. raise HTTPException(status_code=500, detail="Internal server error")
  252. return app