pp_shituv2.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  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. from operator import attrgetter
  16. from typing import Any, Dict, List
  17. from fastapi import FastAPI
  18. from ....pipelines.components import IndexData
  19. from ...infra import utils as serving_utils
  20. from ...infra.config import AppConfig
  21. from ...infra.models import ResultResponse
  22. from ...schemas import pp_shituv2 as schema
  23. from .._app import create_app, primary_operation
  24. from ._common import image_recognition as ir_common
  25. def create_pipeline_app(pipeline: Any, app_config: AppConfig) -> FastAPI:
  26. app, ctx = create_app(
  27. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  28. )
  29. ir_common.update_app_context(ctx)
  30. @primary_operation(
  31. app,
  32. schema.BUILD_INDEX_ENDPOINT,
  33. "buildIndex",
  34. )
  35. async def _build_index(
  36. request: schema.BuildIndexRequest,
  37. ) -> ResultResponse[schema.BuildIndexResult]:
  38. pipeline = ctx.pipeline
  39. aiohttp_session = ctx.aiohttp_session
  40. file_bytes_list = await asyncio.gather(
  41. *(
  42. serving_utils.get_raw_bytes_async(img, aiohttp_session)
  43. for img in map(attrgetter("image"), request.imageLabelPairs)
  44. )
  45. )
  46. images = [serving_utils.image_bytes_to_array(item) for item in file_bytes_list]
  47. labels = [pair.label for pair in request.imageLabelPairs]
  48. # TODO: Support specifying `index_type` and `metric_type` in the
  49. # request
  50. index_data = await pipeline.call(
  51. pipeline.pipeline.build_index,
  52. images,
  53. labels,
  54. index_type="Flat",
  55. metric_type="IP",
  56. )
  57. index_storage = ctx.extra["index_storage"]
  58. index_key = ir_common.generate_index_key()
  59. index_data_bytes = index_data.to_bytes()
  60. await serving_utils.call_async(index_storage.set, index_key, index_data_bytes)
  61. return ResultResponse[schema.BuildIndexResult](
  62. logId=serving_utils.generate_log_id(),
  63. result=schema.BuildIndexResult(indexKey=index_key, idMap=index_data.id_map),
  64. )
  65. @primary_operation(
  66. app,
  67. schema.ADD_IMAGES_TO_INDEX_ENDPOINT,
  68. "addImagesToIndex",
  69. )
  70. async def _add_images_to_index(
  71. request: schema.AddImagesToIndexRequest,
  72. ) -> ResultResponse[schema.AddImagesToIndexResult]:
  73. pipeline = ctx.pipeline
  74. aiohttp_session = ctx.aiohttp_session
  75. file_bytes_list = await asyncio.gather(
  76. *(
  77. serving_utils.get_raw_bytes_async(img, aiohttp_session)
  78. for img in map(attrgetter("image"), request.imageLabelPairs)
  79. )
  80. )
  81. images = [serving_utils.image_bytes_to_array(item) for item in file_bytes_list]
  82. labels = [pair.label for pair in request.imageLabelPairs]
  83. index_storage = ctx.extra["index_storage"]
  84. index_data_bytes = await serving_utils.call_async(
  85. index_storage.get, request.indexKey
  86. )
  87. index_data = IndexData.from_bytes(index_data_bytes)
  88. index_data = await pipeline.call(
  89. pipeline.pipeline.append_index, images, labels, index_data
  90. )
  91. index_data_bytes = index_data.to_bytes()
  92. await serving_utils.call_async(
  93. index_storage.set, request.indexKey, index_data_bytes
  94. )
  95. return ResultResponse[schema.AddImagesToIndexResult](
  96. logId=serving_utils.generate_log_id(),
  97. result=schema.AddImagesToIndexResult(idMap=index_data.id_map),
  98. )
  99. @primary_operation(
  100. app,
  101. schema.REMOVE_IMAGES_FROM_INDEX_ENDPOINT,
  102. "removeImagesFromIndex",
  103. )
  104. async def _remove_images_from_index(
  105. request: schema.RemoveImagesFromIndexRequest,
  106. ) -> ResultResponse[schema.RemoveImagesFromIndexResult]:
  107. pipeline = ctx.pipeline
  108. index_storage = ctx.extra["index_storage"]
  109. index_data_bytes = await serving_utils.call_async(
  110. index_storage.get, request.indexKey
  111. )
  112. index_data = IndexData.from_bytes(index_data_bytes)
  113. index_data = await pipeline.call(
  114. pipeline.pipeline.remove_index, request.ids, index_data
  115. )
  116. index_data_bytes = index_data.to_bytes()
  117. await serving_utils.call_async(
  118. index_storage.set, request.indexKey, index_data_bytes
  119. )
  120. return ResultResponse[schema.RemoveImagesFromIndexResult](
  121. logId=serving_utils.generate_log_id(),
  122. result=schema.RemoveImagesFromIndexResult(idMap=index_data.id_map),
  123. )
  124. @primary_operation(
  125. app,
  126. schema.INFER_ENDPOINT,
  127. "infer",
  128. )
  129. async def _infer(
  130. request: schema.InferRequest,
  131. ) -> ResultResponse[schema.InferResult]:
  132. pipeline = ctx.pipeline
  133. aiohttp_session = ctx.aiohttp_session
  134. image_bytes = await serving_utils.get_raw_bytes_async(
  135. request.image, aiohttp_session
  136. )
  137. image = serving_utils.image_bytes_to_array(image_bytes)
  138. if request.indexKey is not None:
  139. index_storage = ctx.extra["index_storage"]
  140. index_data_bytes = await serving_utils.call_async(
  141. index_storage.get, request.indexKey
  142. )
  143. index_data = IndexData.from_bytes(index_data_bytes)
  144. else:
  145. index_data = None
  146. result = list(
  147. await pipeline.call(
  148. pipeline.pipeline.predict,
  149. image,
  150. index=index_data,
  151. det_threshold=request.detThreshold,
  152. rec_threshold=request.recThreshold,
  153. hamming_radius=request.hammingRadius,
  154. topk=request.topk,
  155. )
  156. )[0]
  157. objs: List[Dict[str, Any]] = []
  158. for obj in result["boxes"]:
  159. rec_results: List[Dict[str, Any]] = []
  160. if obj["rec_scores"] is not None:
  161. for label, score in zip(obj["labels"], obj["rec_scores"]):
  162. rec_results.append(
  163. dict(
  164. label=label,
  165. score=score,
  166. )
  167. )
  168. objs.append(
  169. dict(
  170. bbox=obj["coordinate"],
  171. recResults=rec_results,
  172. score=obj["det_score"],
  173. )
  174. )
  175. if ctx.config.visualize:
  176. output_image_base64 = serving_utils.base64_encode(
  177. serving_utils.image_to_bytes(result.img["res"])
  178. )
  179. else:
  180. output_image_base64 = None
  181. return ResultResponse[schema.InferResult](
  182. logId=serving_utils.generate_log_id(),
  183. result=schema.InferResult(detectedObjects=objs, image=output_image_base64),
  184. )
  185. return app