predictor.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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 shutil
  15. import tempfile
  16. from typing import Any, Dict, Iterator, List, Tuple
  17. import lazy_paddle
  18. from ....modules.m_3d_bev_detection.model_list import MODELS
  19. from ....utils import logging
  20. from ....utils.func_register import FuncRegister
  21. from ...common.batch_sampler import Det3DBatchSampler
  22. from ...common.reader import ReadNuscenesData
  23. from ..base import BasePredictor
  24. from ..base.predictor.base_predictor import PredictionWrap
  25. from .processors import (
  26. GetInferInput,
  27. LoadMultiViewImageFromFiles,
  28. LoadPointsFromFile,
  29. LoadPointsFromMultiSweeps,
  30. NormalizeImage,
  31. PadImage,
  32. ResizeImage,
  33. SampleFilterByKey,
  34. )
  35. from .result import BEV3DDetResult
  36. class BEVDet3DPredictor(BasePredictor):
  37. """BEVDet3DPredictor that inherits from BasePredictor."""
  38. entities = MODELS
  39. _FUNC_MAP = {}
  40. register = FuncRegister(_FUNC_MAP)
  41. def __init__(self, *args: List, **kwargs: Dict) -> None:
  42. """Initializes BEVDet3DPredictor.
  43. Args:
  44. *args: Arbitrary positional arguments passed to the superclass.
  45. **kwargs: Arbitrary keyword arguments passed to the superclass.
  46. """
  47. self.temp_dir = tempfile.mkdtemp()
  48. logging.info(
  49. f"infer data will be stored in temporary directory {self.temp_dir}"
  50. )
  51. super().__init__(*args, **kwargs)
  52. self.pre_tfs, self.infer = self._build()
  53. def _build_batch_sampler(self) -> Det3DBatchSampler:
  54. """Builds and returns an Det3DBatchSampler instance.
  55. Returns:
  56. Det3DBatchSampler: An instance of Det3DBatchSampler.
  57. """
  58. return Det3DBatchSampler(temp_dir=self.temp_dir)
  59. def _get_result_class(self) -> type:
  60. """Returns the result class, BEV3DDetResult.
  61. Returns:
  62. type: The BEV3DDetResult class.
  63. """
  64. return BEV3DDetResult
  65. def _build(self) -> Tuple:
  66. """Build the preprocessors and inference engine based on the configuration.
  67. Returns:
  68. tuple: A tuple containing the preprocessors and inference engine.
  69. """
  70. if (
  71. lazy_paddle.is_compiled_with_cuda()
  72. and not lazy_paddle.is_compiled_with_rocm()
  73. ):
  74. from ....ops.iou3d_nms import nms_gpu # noqa: F401
  75. from ....ops.voxelize import hard_voxelize # noqa: F401
  76. else:
  77. logging.error("3D BEVFusion custom ops only support GPU platform!")
  78. pre_tfs = {"Read": ReadNuscenesData()}
  79. for cfg in self.config["PreProcess"]["transform_ops"]:
  80. tf_key = list(cfg.keys())[0]
  81. func = self._FUNC_MAP[tf_key]
  82. args = cfg.get(tf_key, {})
  83. name, op = func(self, **args) if args else func(self)
  84. if op:
  85. pre_tfs[name] = op
  86. pre_tfs["GetInferInput"] = GetInferInput()
  87. infer = self.create_static_infer()
  88. return pre_tfs, infer
  89. def _format_output(
  90. self, infer_input: List[Any], outs: List[Any], img_metas: Dict[str, Any]
  91. ) -> Dict[str, Any]:
  92. """format inference input and output into predict result
  93. Args:
  94. infer_input(List): Model infer inputs with list containing images, points and lidar2img matrix.
  95. outs(List): Model infer output containing bboxes, scores, labels result.
  96. img_metas(Dict): Image metas info of input sample.
  97. Returns:
  98. Dict: A Dict containing formatted inference output results.
  99. """
  100. input_lidar_path = img_metas["input_lidar_path"]
  101. input_img_paths = img_metas["input_img_paths"]
  102. sample_id = img_metas["sample_id"]
  103. results = {}
  104. out_bboxes_3d = []
  105. out_scores_3d = []
  106. out_labels_3d = []
  107. input_imgs = []
  108. input_points = []
  109. input_lidar2imgs = []
  110. input_ids = []
  111. input_lidar_path_list = []
  112. input_img_paths_list = []
  113. out_bboxes_3d.append(outs[0])
  114. out_labels_3d.append(outs[1])
  115. out_scores_3d.append(outs[2])
  116. input_imgs.append(infer_input[1])
  117. input_points.append(infer_input[0])
  118. input_lidar2imgs.append(infer_input[2])
  119. input_ids.append(sample_id)
  120. input_lidar_path_list.append(input_lidar_path)
  121. input_img_paths_list.append(input_img_paths)
  122. results["input_path"] = input_lidar_path_list
  123. results["input_img_paths"] = input_img_paths_list
  124. results["sample_id"] = input_ids
  125. results["boxes_3d"] = out_bboxes_3d
  126. results["labels_3d"] = out_labels_3d
  127. results["scores_3d"] = out_scores_3d
  128. return results
  129. def process(self, batch_data: List[str]) -> Dict[str, Any]:
  130. """
  131. Process a batch of data through the preprocessing and inference.
  132. Args:
  133. batch_data (List[str]): A batch of input data (e.g., sample anno file paths).
  134. Returns:
  135. dict: A dictionary containing the input path, input img, input points, input lidar2img, output bboxes, output labels, output scores and label names. Keys include 'input_path', 'input_img', 'input_points', 'input_lidar2img', 'boxes_3d', 'labels_3d' and 'scores_3d'.
  136. """
  137. sample = self.pre_tfs["Read"](batch_data=batch_data)
  138. sample = self.pre_tfs["LoadPointsFromFile"](results=sample[0])
  139. sample = self.pre_tfs["LoadPointsFromMultiSweeps"](results=sample)
  140. sample = self.pre_tfs["LoadMultiViewImageFromFiles"](sample=sample)
  141. sample = self.pre_tfs["ResizeImage"](results=sample)
  142. sample = self.pre_tfs["NormalizeImage"](results=sample)
  143. sample = self.pre_tfs["PadImage"](results=sample)
  144. sample = self.pre_tfs["SampleFilterByKey"](sample=sample)
  145. infer_input, img_metas = self.pre_tfs["GetInferInput"](sample=sample)
  146. infer_output = self.infer(x=infer_input)
  147. results = self._format_output(infer_input, infer_output, img_metas)
  148. return results
  149. @register("LoadPointsFromFile")
  150. def build_load_img_from_file(
  151. self, load_dim=6, use_dim=[0, 1, 2], shift_height=False, use_color=False
  152. ):
  153. return "LoadPointsFromFile", LoadPointsFromFile(
  154. load_dim=load_dim,
  155. use_dim=use_dim,
  156. shift_height=shift_height,
  157. use_color=use_color,
  158. )
  159. @register("LoadPointsFromMultiSweeps")
  160. def build_load_points_from_multi_sweeps(
  161. self,
  162. sweeps_num=10,
  163. load_dim=5,
  164. use_dim=[0, 1, 2, 4],
  165. pad_empty_sweeps=False,
  166. remove_close=False,
  167. test_mode=False,
  168. point_cloud_angle_range=None,
  169. ):
  170. return "LoadPointsFromMultiSweeps", LoadPointsFromMultiSweeps(
  171. sweeps_num=sweeps_num,
  172. load_dim=load_dim,
  173. use_dim=use_dim,
  174. pad_empty_sweeps=pad_empty_sweeps,
  175. remove_close=remove_close,
  176. test_mode=test_mode,
  177. point_cloud_angle_range=point_cloud_angle_range,
  178. )
  179. @register("LoadMultiViewImageFromFiles")
  180. def build_load_multi_view_image_from_files(
  181. self,
  182. to_float32=False,
  183. project_pts_to_img_depth=False,
  184. cam_depth_range=[4.0, 45.0, 1.0],
  185. constant_std=0.5,
  186. imread_flag=-1,
  187. ):
  188. return "LoadMultiViewImageFromFiles", LoadMultiViewImageFromFiles(
  189. to_float32=to_float32,
  190. project_pts_to_img_depth=project_pts_to_img_depth,
  191. cam_depth_range=cam_depth_range,
  192. constant_std=constant_std,
  193. imread_flag=imread_flag,
  194. )
  195. @register("ResizeImage")
  196. def build_resize_image(
  197. self,
  198. img_scale=None,
  199. multiscale_mode="range",
  200. ratio_range=None,
  201. keep_ratio=True,
  202. bbox_clip_border=True,
  203. backend="cv2",
  204. override=False,
  205. ):
  206. return "ResizeImage", ResizeImage(
  207. img_scale=img_scale,
  208. multiscale_mode=multiscale_mode,
  209. ratio_range=ratio_range,
  210. keep_ratio=keep_ratio,
  211. bbox_clip_border=bbox_clip_border,
  212. backend=backend,
  213. override=override,
  214. )
  215. @register("NormalizeImage")
  216. def build_normalize_image(self, mean, std, to_rgb=True):
  217. return "NormalizeImage", NormalizeImage(mean=mean, std=std, to_rgb=to_rgb)
  218. @register("PadImage")
  219. def build_pad_image(self, size=None, size_divisor=None, pad_val=0):
  220. return "PadImage", PadImage(
  221. size=size, size_divisor=size_divisor, pad_val=pad_val
  222. )
  223. @register("SampleFilterByKey")
  224. def build_sample_filter_by_key(
  225. self,
  226. keys,
  227. meta_keys=(
  228. "filename",
  229. "ori_shape",
  230. "img_shape",
  231. "lidar2img",
  232. "depth2img",
  233. "cam2img",
  234. "pad_shape",
  235. "scale_factor",
  236. "flip",
  237. "pcd_horizontal_flip",
  238. "pcd_vertical_flip",
  239. "box_type_3d",
  240. "img_norm_cfg",
  241. "pcd_trans",
  242. "sample_idx",
  243. "pcd_scale_factor",
  244. "pcd_rotation",
  245. "pts_filename",
  246. "transformation_3d_flow",
  247. ),
  248. ):
  249. return "SampleFilterByKey", SampleFilterByKey(keys=keys, meta_keys=meta_keys)
  250. @register("GetInferInput")
  251. def build_get_infer_input(self):
  252. return "GetInferInput", GetInferInput()
  253. def apply(self, input: Any, **kwargs) -> Iterator[Any]:
  254. """
  255. Do predicting with the input data and yields predictions.
  256. Args:
  257. input (Any): The input data to be predicted.
  258. Yields:
  259. Iterator[Any]: An iterator yielding prediction results.
  260. """
  261. try:
  262. for batch_data in self.batch_sampler(input):
  263. prediction = self.process(batch_data, **kwargs)
  264. prediction = PredictionWrap(prediction, len(batch_data))
  265. for idx in range(len(batch_data)):
  266. yield self.result_class(prediction.get_by_idx(idx))
  267. except Exception as e:
  268. raise e
  269. finally:
  270. shutil.rmtree(self.temp_dir)