pp_shitu_v2.py 10 KB

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