processors.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983
  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 numbers
  15. import cv2
  16. import numpy as np
  17. from ...common.reader.det_3d_reader import Sample
  18. from ...utils.benchmark import benchmark
  19. cv2_interp_codes = {
  20. "nearest": cv2.INTER_NEAREST,
  21. "bilinear": cv2.INTER_LINEAR,
  22. "bicubic": cv2.INTER_CUBIC,
  23. "area": cv2.INTER_AREA,
  24. "lanczos": cv2.INTER_LANCZOS4,
  25. }
  26. @benchmark.timeit
  27. class LoadPointsFromFile:
  28. """Load points from a file and process them according to specified parameters."""
  29. def __init__(
  30. self, load_dim=6, use_dim=[0, 1, 2], shift_height=False, use_color=False
  31. ):
  32. """Initializes the LoadPointsFromFile object.
  33. Args:
  34. load_dim (int): Dimensions loaded in points.
  35. use_dim (list or int): Dimensions used in points. If int, will use a range from 0 to use_dim (exclusive).
  36. shift_height (bool): Whether to shift height values.
  37. use_color (bool): Whether to include color attributes in the loaded points.
  38. """
  39. self.shift_height = shift_height
  40. self.use_color = use_color
  41. if isinstance(use_dim, int):
  42. use_dim = list(range(use_dim))
  43. assert (
  44. max(use_dim) < load_dim
  45. ), f"Expect all used dimensions < {load_dim}, got {use_dim}"
  46. self.load_dim = load_dim
  47. self.use_dim = use_dim
  48. def _load_points(self, pts_filename):
  49. """Private function to load point clouds data from a file.
  50. Args:
  51. pts_filename (str): Path to the point cloud file.
  52. Returns:
  53. numpy.ndarray: Loaded point cloud data.
  54. """
  55. points = np.fromfile(pts_filename, dtype=np.float32)
  56. return points
  57. def __call__(self, results):
  58. """Call function to load points data from file and process it.
  59. Args:
  60. results (dict): Dictionary containing the 'pts_filename' key with the path to the point cloud file.
  61. Returns:
  62. dict: Updated results dictionary with 'points' key added.
  63. """
  64. pts_filename = results["pts_filename"]
  65. points = self._load_points(pts_filename)
  66. points = points.reshape(-1, self.load_dim)
  67. points = points[:, self.use_dim]
  68. attribute_dims = None
  69. if self.shift_height:
  70. floor_height = np.percentile(points[:, 2], 0.99)
  71. height = points[:, 2] - floor_height
  72. points = np.concatenate(
  73. [points[:, :3], np.expand_dims(height, 1), points[:, 3:]], 1
  74. )
  75. attribute_dims = dict(height=3)
  76. if self.use_color:
  77. assert len(self.use_dim) >= 6
  78. if attribute_dims is None:
  79. attribute_dims = dict()
  80. attribute_dims.update(
  81. dict(
  82. color=[
  83. points.shape[1] - 3,
  84. points.shape[1] - 2,
  85. points.shape[1] - 1,
  86. ]
  87. )
  88. )
  89. results["points"] = points
  90. return results
  91. @benchmark.timeit
  92. class LoadPointsFromMultiSweeps(object):
  93. """Load points from multiple sweeps.This is usually used for nuScenes dataset to utilize previous sweeps."""
  94. def __init__(
  95. self,
  96. sweeps_num=10,
  97. load_dim=5,
  98. use_dim=[0, 1, 2, 4],
  99. pad_empty_sweeps=False,
  100. remove_close=False,
  101. test_mode=False,
  102. point_cloud_angle_range=None,
  103. ):
  104. """Initializes the LoadPointsFromMultiSweeps object
  105. Args:
  106. sweeps_num (int): Number of sweeps. Defaults to 10.
  107. load_dim (int): Dimension number of the loaded points. Defaults to 5.
  108. use_dim (list[int]): Which dimension to use. Defaults to [0, 1, 2, 4].
  109. for more details. Defaults to dict(backend='disk').
  110. pad_empty_sweeps (bool): Whether to repeat keyframe when
  111. sweeps is empty. Defaults to False.
  112. remove_close (bool): Whether to remove close points.
  113. Defaults to False.
  114. test_mode (bool): If test_model=True used for testing, it will not
  115. randomly sample sweeps but select the nearest N frames.
  116. Defaults to False.
  117. """
  118. self.load_dim = load_dim
  119. self.sweeps_num = sweeps_num
  120. self.use_dim = use_dim
  121. self.pad_empty_sweeps = pad_empty_sweeps
  122. self.remove_close = remove_close
  123. self.test_mode = test_mode
  124. if point_cloud_angle_range is not None:
  125. self.filter_by_angle = True
  126. self.point_cloud_angle_range = point_cloud_angle_range
  127. print(point_cloud_angle_range)
  128. else:
  129. self.filter_by_angle = False
  130. # self.point_cloud_angle_range = point_cloud_angle_range
  131. def _load_points(self, pts_filename):
  132. """Private function to load point clouds data.
  133. Args:
  134. pts_filename (str): Filename of point clouds data.
  135. Returns:
  136. np.ndarray: An array containing point clouds data.
  137. """
  138. points = np.fromfile(pts_filename, dtype=np.float32)
  139. return points
  140. def _remove_close(self, points, radius=1.0):
  141. """Removes point too close within a certain radius from origin.
  142. Args:
  143. points (np.ndarray): Sweep points.
  144. radius (float): Radius below which points are removed.
  145. Defaults to 1.0.
  146. Returns:
  147. np.ndarray: Points after removing.
  148. """
  149. if isinstance(points, np.ndarray):
  150. points_numpy = points
  151. else:
  152. raise NotImplementedError
  153. x_filt = np.abs(points_numpy[:, 0]) < radius
  154. y_filt = np.abs(points_numpy[:, 1]) < radius
  155. not_close = np.logical_not(np.logical_and(x_filt, y_filt))
  156. return points[not_close]
  157. def filter_point_by_angle(self, points):
  158. """
  159. Filters points based on their angle in relation to the origin.
  160. Args:
  161. points (np.ndarray): An array of points with shape (N, 2), where each row
  162. is a point in 2D space.
  163. Returns:
  164. np.ndarray: A filtered array of points that fall within the specified
  165. angle range.
  166. """
  167. if isinstance(points, np.ndarray):
  168. points_numpy = points
  169. else:
  170. raise NotImplementedError
  171. pts_phi = (
  172. np.arctan(points_numpy[:, 0] / points_numpy[:, 1])
  173. + (points_numpy[:, 1] < 0) * np.pi
  174. + np.pi * 2
  175. ) % (np.pi * 2)
  176. pts_phi[pts_phi > np.pi] -= np.pi * 2
  177. pts_phi = pts_phi / np.pi * 180
  178. assert np.all(-180 <= pts_phi) and np.all(pts_phi <= 180)
  179. filt = np.logical_and(
  180. pts_phi >= self.point_cloud_angle_range[0],
  181. pts_phi <= self.point_cloud_angle_range[1],
  182. )
  183. return points[filt]
  184. def __call__(self, results):
  185. """Call function to load multi-sweep point clouds from files.
  186. Args:
  187. results (dict): Result dict containing multi-sweep point cloud \
  188. filenames.
  189. Returns:
  190. dict: The result dict containing the multi-sweep points data. \
  191. Added key and value are described below.
  192. - points (np.ndarray): Multi-sweep point cloud arrays.
  193. """
  194. points = results["points"]
  195. points[:, 4] = 0
  196. sweep_points_list = [points]
  197. ts = results["timestamp"]
  198. if self.pad_empty_sweeps and len(results["sweeps"]) == 0:
  199. for i in range(self.sweeps_num):
  200. if self.remove_close:
  201. sweep_points_list.append(self._remove_close(points))
  202. else:
  203. sweep_points_list.append(points)
  204. else:
  205. if len(results["sweeps"]) <= self.sweeps_num:
  206. choices = np.arange(len(results["sweeps"]))
  207. elif self.test_mode:
  208. choices = np.arange(self.sweeps_num)
  209. else:
  210. choices = np.random.choice(
  211. len(results["sweeps"]), self.sweeps_num, replace=False
  212. )
  213. for idx in choices:
  214. sweep = results["sweeps"][idx]
  215. points_sweep = self._load_points(sweep["data_path"])
  216. points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
  217. if self.remove_close:
  218. points_sweep = self._remove_close(points_sweep)
  219. sweep_ts = sweep["timestamp"] / 1e6
  220. points_sweep[:, :3] = (
  221. points_sweep[:, :3] @ sweep["sensor2lidar_rotation"].T
  222. )
  223. points_sweep[:, :3] += sweep["sensor2lidar_translation"]
  224. points_sweep[:, 4] = ts - sweep_ts
  225. # points_sweep = points.new_point(points_sweep)
  226. sweep_points_list.append(points_sweep)
  227. points = np.concatenate(sweep_points_list, axis=0)
  228. if self.filter_by_angle:
  229. points = self.filter_point_by_angle(points)
  230. points = points[:, self.use_dim]
  231. results["points"] = points
  232. return results
  233. @benchmark.timeit
  234. class LoadMultiViewImageFromFiles:
  235. """Load multi-view images from files."""
  236. def __init__(
  237. self,
  238. to_float32=False,
  239. project_pts_to_img_depth=False,
  240. cam_depth_range=[4.0, 45.0, 1.0],
  241. constant_std=0.5,
  242. imread_flag=-1,
  243. ):
  244. """
  245. Initializes the LoadMultiViewImageFromFiles object.
  246. Args:
  247. to_float32 (bool): Whether to convert the loaded images to float32. Default: False.
  248. project_pts_to_img_depth (bool): Whether to project points to image depth. Default: False.
  249. cam_depth_range (list): Camera depth range in the format [min, max, focal]. Default: [4.0, 45.0, 1.0].
  250. constant_std (float): Constant standard deviation for normalization. Default: 0.5.
  251. imread_flag (int): Flag determining the color type of the loaded image.
  252. - -1: cv2.IMREAD_UNCHANGED
  253. - 0: cv2.IMREAD_GRAYSCALE
  254. - 1: cv2.IMREAD_COLOR
  255. Default: -1.
  256. """
  257. self.to_float32 = to_float32
  258. self.project_pts_to_img_depth = project_pts_to_img_depth
  259. self.cam_depth_range = cam_depth_range
  260. self.constant_std = constant_std
  261. self.imread_flag = imread_flag
  262. def __call__(self, sample):
  263. """
  264. Call method to load multi-view image from files and update the sample dictionary.
  265. Args:
  266. sample (dict): Dictionary containing the image filename key.
  267. Returns:
  268. dict: Updated sample dictionary with loaded images and additional information.
  269. """
  270. filename = sample["img_filename"]
  271. img = np.stack(
  272. [cv2.imread(name, self.imread_flag) for name in filename], axis=-1
  273. )
  274. if self.to_float32:
  275. img = img.astype(np.float32)
  276. sample["filename"] = filename
  277. sample["img"] = [img[..., i] for i in range(img.shape[-1])]
  278. sample["img_shape"] = img.shape
  279. sample["ori_shape"] = img.shape
  280. sample["pad_shape"] = img.shape
  281. # sample['scale_factor'] = 1.0
  282. num_channels = 1 if len(img.shape) < 3 else img.shape[2]
  283. sample["img_norm_cfg"] = dict(
  284. mean=np.zeros(num_channels, dtype=np.float32),
  285. std=np.ones(num_channels, dtype=np.float32),
  286. to_rgb=False,
  287. )
  288. sample["img_fields"] = ["img"]
  289. return sample
  290. @benchmark.timeit
  291. class ResizeImage:
  292. """Resize images & bbox & mask."""
  293. def __init__(
  294. self,
  295. img_scale=None,
  296. multiscale_mode="range",
  297. ratio_range=None,
  298. keep_ratio=True,
  299. bbox_clip_border=True,
  300. backend="cv2",
  301. override=False,
  302. ):
  303. """Initializes the ResizeImage object.
  304. Args:
  305. img_scale (list or int, optional): The scale of the image. If a single integer is provided, it will be converted to a list. Defaults to None.
  306. multiscale_mode (str): The mode for multiscale resizing. Can be "value" or "range". Defaults to "range".
  307. ratio_range (list, optional): The range of image aspect ratios. Only used when img_scale is a single value. Defaults to None.
  308. keep_ratio (bool): Whether to keep the aspect ratio when resizing. Defaults to True.
  309. bbox_clip_border (bool): Whether to clip the bounding box to the image border. Defaults to True.
  310. backend (str): The backend to use for image resizing. Can be "cv2". Defaults to "cv2".
  311. override (bool): Whether to override certain resize parameters. Note: This option needs refactoring. Defaults to False.
  312. """
  313. if img_scale is None:
  314. self.img_scale = None
  315. else:
  316. if isinstance(img_scale, list):
  317. self.img_scale = img_scale
  318. else:
  319. self.img_scale = [img_scale]
  320. if ratio_range is not None:
  321. # mode 1: given a scale and a range of image ratio
  322. assert len(self.img_scale) == 1
  323. else:
  324. # mode 2: given multiple scales or a range of scales
  325. assert multiscale_mode in ["value", "range"]
  326. self.backend = backend
  327. self.multiscale_mode = multiscale_mode
  328. self.ratio_range = ratio_range
  329. self.keep_ratio = keep_ratio
  330. # TODO: refactor the override option in Resize
  331. self.override = override
  332. self.bbox_clip_border = bbox_clip_border
  333. @staticmethod
  334. def random_select(img_scales):
  335. """Randomly select an img_scale from the given list of candidates.
  336. Args:
  337. img_scales (list): A list of image scales to choose from.
  338. Returns:
  339. tuple: A tuple containing the selected image scale and its index in the list.
  340. """
  341. scale_idx = np.random.randint(len(img_scales))
  342. img_scale = img_scales[scale_idx]
  343. return img_scale, scale_idx
  344. @staticmethod
  345. def random_sample(img_scales):
  346. """
  347. Randomly sample an img_scale when `multiscale_mode` is set to 'range'.
  348. Args:
  349. img_scales (list of tuples): A list of tuples, where each tuple contains
  350. the minimum and maximum scale dimensions for an image.
  351. Returns:
  352. tuple: A tuple containing the randomly sampled img_scale (long_edge, short_edge)
  353. and None (to maintain function signature compatibility).
  354. """
  355. img_scale_long = [max(s) for s in img_scales]
  356. img_scale_short = [min(s) for s in img_scales]
  357. long_edge = np.random.randint(min(img_scale_long), max(img_scale_long) + 1)
  358. short_edge = np.random.randint(min(img_scale_short), max(img_scale_short) + 1)
  359. img_scale = (long_edge, short_edge)
  360. return img_scale, None
  361. @staticmethod
  362. def random_sample_ratio(img_scale, ratio_range):
  363. """
  364. Randomly sample an img_scale based on the specified ratio_range.
  365. Args:
  366. img_scale (list): A list of two integers representing the minimum and maximum
  367. scale for the image.
  368. ratio_range (tuple): A tuple of two floats representing the minimum and maximum
  369. ratio for sampling the img_scale.
  370. Returns:
  371. tuple: A tuple containing the sampled scale (as a tuple of two integers)
  372. and None.
  373. """
  374. assert isinstance(img_scale, list) and len(img_scale) == 2
  375. min_ratio, max_ratio = ratio_range
  376. assert min_ratio <= max_ratio
  377. ratio = np.random.random_sample() * (max_ratio - min_ratio) + min_ratio
  378. scale = int(img_scale[0] * ratio), int(img_scale[1] * ratio)
  379. return scale, None
  380. def _random_scale(self, results):
  381. """Randomly sample an img_scale according to `ratio_range` and `multiscale_mode`.
  382. Args:
  383. results (dict): A dictionary to store the sampled scale and its index.
  384. Returns:
  385. None. The sampled scale and its index are stored in `results` dictionary.
  386. """
  387. if self.ratio_range is not None:
  388. scale, scale_idx = self.random_sample_ratio(
  389. self.img_scale[0], self.ratio_range
  390. )
  391. elif len(self.img_scale) == 1:
  392. scale, scale_idx = self.img_scale[0], 0
  393. elif self.multiscale_mode == "range":
  394. scale, scale_idx = self.random_sample(self.img_scale)
  395. elif self.multiscale_mode == "value":
  396. scale, scale_idx = self.random_select(self.img_scale)
  397. else:
  398. raise NotImplementedError
  399. results["scale"] = scale
  400. results["scale_idx"] = scale_idx
  401. def _resize_img(self, results):
  402. """Resize images based on the scale factor provided in ``results['scale']`` while maintaining the aspect ratio if ``self.keep_ratio`` is True.
  403. Args:
  404. results (dict): A dictionary containing image fields and their corresponding scales.
  405. Returns:
  406. None. The ``results`` dictionary is modified in place with resized images and additional fields like `img_shape`, `pad_shape`, `scale_factor`, and `keep_ratio`.
  407. """
  408. for key in results.get("img_fields", ["img"]):
  409. for idx in range(len(results["img"])):
  410. if self.keep_ratio:
  411. img, scale_factor = self.imrescale(
  412. results[key][idx],
  413. results["scale"],
  414. interpolation="bilinear" if key == "img" else "nearest",
  415. return_scale=True,
  416. backend=self.backend,
  417. )
  418. new_h, new_w = img.shape[:2]
  419. h, w = results[key][idx].shape[:2]
  420. w_scale = new_w / w
  421. h_scale = new_h / h
  422. else:
  423. raise NotImplementedError
  424. results[key][idx] = img
  425. scale_factor = np.array(
  426. [w_scale, h_scale, w_scale, h_scale], dtype=np.float32
  427. )
  428. results["img_shape"] = img.shape
  429. # in case that there is no padding
  430. results["pad_shape"] = img.shape
  431. results["scale_factor"] = scale_factor
  432. results["keep_ratio"] = self.keep_ratio
  433. def rescale_size(self, old_size, scale, return_scale=False):
  434. """
  435. Calculate the new size to be rescaled to based on the given scale.
  436. Args:
  437. old_size (tuple): A tuple containing the width and height of the original size.
  438. scale (float, int, or list of int): The scale factor or a list of integers representing the maximum and minimum allowed size.
  439. return_scale (bool): Whether to return the scale factor along with the new size.
  440. Returns:
  441. tuple: A tuple containing the new size and optionally the scale factor if return_scale is True.
  442. """
  443. w, h = old_size
  444. if isinstance(scale, (float, int)):
  445. if scale <= 0:
  446. raise ValueError(f"Invalid scale {scale}, must be positive.")
  447. scale_factor = scale
  448. elif isinstance(scale, list):
  449. max_long_edge = max(scale)
  450. max_short_edge = min(scale)
  451. scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
  452. else:
  453. raise TypeError(
  454. f"Scale must be a number or list of int, but got {type(scale)}"
  455. )
  456. def _scale_size(size, scale):
  457. if isinstance(scale, (float, int)):
  458. scale = (scale, scale)
  459. w, h = size
  460. return int(w * float(scale[0]) + 0.5), int(h * float(scale[1]) + 0.5)
  461. new_size = _scale_size((w, h), scale_factor)
  462. if return_scale:
  463. return new_size, scale_factor
  464. else:
  465. return new_size
  466. def imrescale(
  467. self, img, scale, return_scale=False, interpolation="bilinear", backend=None
  468. ):
  469. """Resize image while keeping the aspect ratio.
  470. Args:
  471. img (numpy.ndarray): The input image.
  472. scale (float): The scaling factor.
  473. return_scale (bool): Whether to return the scaling factor along with the resized image.
  474. interpolation (str): The interpolation method to use. Defaults to 'bilinear'.
  475. backend (str): The backend to use for resizing. Defaults to None.
  476. Returns:
  477. tuple or numpy.ndarray: The resized image, and optionally the scaling factor.
  478. """
  479. h, w = img.shape[:2]
  480. new_size, scale_factor = self.rescale_size((w, h), scale, return_scale=True)
  481. rescaled_img = self.imresize(
  482. img, new_size, interpolation=interpolation, backend=backend
  483. )
  484. if return_scale:
  485. return rescaled_img, scale_factor
  486. else:
  487. return rescaled_img
  488. def imresize(
  489. self,
  490. img,
  491. size,
  492. return_scale=False,
  493. interpolation="bilinear",
  494. out=None,
  495. backend=None,
  496. ):
  497. """Resize an image to a given size.
  498. Args:
  499. img (numpy.ndarray): The input image to be resized.
  500. size (tuple): The new size for the image as (height, width).
  501. return_scale (bool): Whether to return the scaling factors along with the resized image.
  502. interpolation (str): The interpolation method to use. Default is 'bilinear'.
  503. out (numpy.ndarray, optional): Output array. If provided, it must have the same shape and dtype as the output array.
  504. backend (str, optional): The backend to use for resizing. Supported backends are 'cv2' and 'pillow'.
  505. Returns:
  506. numpy.ndarray or tuple: The resized image. If return_scale is True, returns a tuple containing the resized image and the scaling factors (w_scale, h_scale).
  507. """
  508. h, w = img.shape[:2]
  509. if backend not in ["cv2", "pillow"]:
  510. raise ValueError(
  511. f"backend: {backend} is not supported for resize."
  512. f"Supported backends are 'cv2', 'pillow'"
  513. )
  514. if backend == "pillow":
  515. raise NotImplementedError
  516. else:
  517. resized_img = cv2.resize(
  518. img, size, dst=out, interpolation=cv2_interp_codes[interpolation]
  519. )
  520. if not return_scale:
  521. return resized_img
  522. else:
  523. w_scale = size[0] / w
  524. h_scale = size[1] / h
  525. return resized_img, w_scale, h_scale
  526. def _resize_bboxes(self, results):
  527. """Resize bounding boxes with `results['scale_factor']`.
  528. Args:
  529. results (dict): A dictionary containing the bounding boxes and other related information.
  530. """
  531. for key in results.get("bbox_fields", []):
  532. bboxes = results[key] * results["scale_factor"]
  533. if self.bbox_clip_border:
  534. img_shape = results["img_shape"]
  535. bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0, img_shape[1])
  536. bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, img_shape[0])
  537. results[key] = bboxes
  538. def _resize_masks(self, results):
  539. """Resize masks with ``results['scale']``"""
  540. raise NotImplementedError
  541. def _resize_seg(self, results):
  542. """Resize semantic segmentation map with ``results['scale']``."""
  543. raise NotImplementedError
  544. def __call__(self, results):
  545. """Call function to resize images, bounding boxes, masks, and semantic segmentation maps according to the provided scale or scale factor.
  546. Args:
  547. results (dict): A dictionary containing the input data, including 'img', 'scale', and optionally 'scale_factor'.
  548. Returns:
  549. dict: A dictionary with the resized data.
  550. """
  551. if "scale" not in results:
  552. if "scale_factor" in results:
  553. img_shape = results["img"][0].shape[:2]
  554. scale_factor = results["scale_factor"]
  555. assert isinstance(scale_factor, float)
  556. results["scale"] = list(
  557. [int(x * scale_factor) for x in img_shape][::-1]
  558. )
  559. else:
  560. self._random_scale(results)
  561. else:
  562. if not self.override:
  563. assert (
  564. "scale_factor" not in results
  565. ), "scale and scale_factor cannot be both set."
  566. else:
  567. results.pop("scale")
  568. if "scale_factor" in results:
  569. results.pop("scale_factor")
  570. self._random_scale(results)
  571. self._resize_img(results)
  572. self._resize_bboxes(results)
  573. return results
  574. @benchmark.timeit
  575. class NormalizeImage:
  576. """Normalize the image."""
  577. """Normalize an image by subtracting the mean and dividing by the standard deviation.
  578. Args:
  579. mean (list or tuple): Mean values for each channel.
  580. std (list or tuple): Standard deviation values for each channel.
  581. to_rgb (bool): Whether to convert the image from BGR to RGB.
  582. """
  583. def __init__(self, mean, std, to_rgb=True):
  584. """Initializes the NormalizeImage class with mean, std, and to_rgb parameters."""
  585. self.mean = np.array(mean, dtype=np.float32)
  586. self.std = np.array(std, dtype=np.float32)
  587. self.to_rgb = to_rgb
  588. def _imnormalize(self, img, mean, std, to_rgb=True):
  589. """Normalize the given image inplace.
  590. Args:
  591. img (numpy.ndarray): The image to normalize.
  592. mean (numpy.ndarray): Mean values for normalization.
  593. std (numpy.ndarray): Standard deviation values for normalization.
  594. to_rgb (bool): Whether to convert the image from BGR to RGB.
  595. Returns:
  596. numpy.ndarray: The normalized image.
  597. """
  598. img = img.copy().astype(np.float32)
  599. mean = np.float64(mean.reshape(1, -1))
  600. stdinv = 1 / np.float64(std.reshape(1, -1))
  601. if to_rgb:
  602. cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img) # inplace
  603. cv2.subtract(img, mean, img) # inplace
  604. cv2.multiply(img, stdinv, img) # inplace
  605. return img
  606. def __call__(self, results):
  607. """Call method to normalize images in the results dictionary.
  608. Args:
  609. results (dict): A dictionary containing image fields to normalize.
  610. Returns:
  611. dict: The results dictionary with normalized images.
  612. """
  613. for key in results.get("img_fields", ["img"]):
  614. if key == "img_depth":
  615. continue
  616. for idx in range(len(results["img"])):
  617. results[key][idx] = self._imnormalize(
  618. results[key][idx], self.mean, self.std, self.to_rgb
  619. )
  620. results["img_norm_cfg"] = dict(mean=self.mean, std=self.std, to_rgb=self.to_rgb)
  621. return results
  622. @benchmark.timeit
  623. class PadImage(object):
  624. """Pad the image & mask."""
  625. def __init__(self, size=None, size_divisor=None, pad_val=0):
  626. self.size = size
  627. self.size_divisor = size_divisor
  628. self.pad_val = pad_val
  629. # only one of size and size_divisor should be valid
  630. assert size is not None or size_divisor is not None
  631. assert size is None or size_divisor is None
  632. def impad(
  633. self, img, *, shape=None, padding=None, pad_val=0, padding_mode="constant"
  634. ):
  635. """Pad the given image to a certain shape or pad on all sides
  636. Args:
  637. img (numpy.ndarray): The input image to be padded.
  638. shape (tuple, optional): Desired output shape in the form (height, width). One of shape or padding must be specified.
  639. padding (int, tuple, optional): Number of pixels to pad on each side of the image. If a single int is provided this
  640. is used to pad all sides with this value. If a tuple of length 2 is provided this is interpreted as (top_bottom, left_right).
  641. If a tuple of length 4 is provided this is interpreted as (top, right, bottom, left).
  642. pad_val (int, list, optional): Pixel value used for padding. If a list is provided, it must have the same length as the
  643. last dimension of the input image. Defaults to 0.
  644. padding_mode (str, optional): Padding mode to use. One of 'constant', 'edge', 'reflect', 'symmetric'.
  645. Defaults to 'constant'.
  646. Returns:
  647. numpy.ndarray: The padded image.
  648. """
  649. assert (shape is not None) ^ (padding is not None)
  650. if shape is not None:
  651. padding = [0, 0, shape[1] - img.shape[1], shape[0] - img.shape[0]]
  652. # check pad_val
  653. if isinstance(pad_val, list):
  654. assert len(pad_val) == img.shape[-1]
  655. elif not isinstance(pad_val, numbers.Number):
  656. raise TypeError(
  657. "pad_val must be a int or a list. " f"But received {type(pad_val)}"
  658. )
  659. # check padding
  660. if isinstance(padding, list) and len(padding) in [2, 4]:
  661. if len(padding) == 2:
  662. padding = [padding[0], padding[1], padding[0], padding[1]]
  663. elif isinstance(padding, numbers.Number):
  664. padding = [padding, padding, padding, padding]
  665. else:
  666. raise ValueError(
  667. "Padding must be a int or a 2, or 4 element list."
  668. f"But received {padding}"
  669. )
  670. # check padding mode
  671. assert padding_mode in ["constant", "edge", "reflect", "symmetric"]
  672. border_type = {
  673. "constant": cv2.BORDER_CONSTANT,
  674. "edge": cv2.BORDER_REPLICATE,
  675. "reflect": cv2.BORDER_REFLECT_101,
  676. "symmetric": cv2.BORDER_REFLECT,
  677. }
  678. img = cv2.copyMakeBorder(
  679. img,
  680. padding[1],
  681. padding[3],
  682. padding[0],
  683. padding[2],
  684. border_type[padding_mode],
  685. value=pad_val,
  686. )
  687. return img
  688. def impad_to_multiple(self, img, divisor, pad_val=0):
  689. """
  690. Pad an image to ensure each edge length is a multiple of a given number.
  691. Args:
  692. img (numpy.ndarray): The input image.
  693. divisor (int): The number to which each edge length should be a multiple.
  694. pad_val (int, optional): The value to pad the image with. Defaults to 0.
  695. Returns:
  696. numpy.ndarray: The padded image.
  697. """
  698. pad_h = int(np.ceil(img.shape[0] / divisor)) * divisor
  699. pad_w = int(np.ceil(img.shape[1] / divisor)) * divisor
  700. return self.impad(img, shape=(pad_h, pad_w), pad_val=pad_val)
  701. def _pad_img(self, results):
  702. """
  703. Pad images according to ``self.size`` or adjust their shapes to be multiples of ``self.size_divisor``.
  704. Args:
  705. results (dict): A dictionary containing image data, with 'img_fields' as an optional key
  706. pointing to a list of image field names.
  707. """
  708. for key in results.get("img_fields", ["img"]):
  709. if self.size is not None:
  710. padded_img = self.impad(
  711. results[key], shape=self.size, pad_val=self.pad_val
  712. )
  713. elif self.size_divisor is not None:
  714. for idx in range(len(results[key])):
  715. padded_img = self.impad_to_multiple(
  716. results[key][idx], self.size_divisor, pad_val=self.pad_val
  717. )
  718. results[key][idx] = padded_img
  719. results["pad_shape"] = padded_img.shape
  720. results["pad_fixed_size"] = self.size
  721. results["pad_size_divisor"] = self.size_divisor
  722. def _pad_masks(self, results):
  723. """Pad masks according to ``results['pad_shape']``."""
  724. raise NotImplementedError
  725. def _pad_seg(self, results):
  726. """Pad semantic segmentation map according to ``results['pad_shape']``."""
  727. raise NotImplementedError
  728. def __call__(self, results):
  729. """Call function to pad images, masks, semantic segmentation maps."""
  730. self._pad_img(results)
  731. return results
  732. @benchmark.timeit
  733. class SampleFilterByKey:
  734. """Collect data from the loader relevant to the specific task."""
  735. def __init__(
  736. self,
  737. keys,
  738. meta_keys=(
  739. "filename",
  740. "ori_shape",
  741. "img_shape",
  742. "lidar2img",
  743. "depth2img",
  744. "cam2img",
  745. "pad_shape",
  746. "scale_factor",
  747. "flip",
  748. "pcd_horizontal_flip",
  749. "pcd_vertical_flip",
  750. "box_type_3d",
  751. "img_norm_cfg",
  752. "pcd_trans",
  753. "sample_idx",
  754. "pcd_scale_factor",
  755. "pcd_rotation",
  756. "pts_filename",
  757. "transformation_3d_flow",
  758. ),
  759. ):
  760. self.keys = keys
  761. self.meta_keys = meta_keys
  762. def __call__(self, sample):
  763. """Call function to filter sample by keys. The keys in `meta_keys` are used to filter metadata from the input sample.
  764. Args:
  765. sample (Sample): The input sample to be filtered.
  766. Returns:
  767. Sample: A new Sample object containing only the filtered metadata and specified keys.
  768. """
  769. filtered_sample = Sample(path=sample.path, modality=sample.modality)
  770. filtered_sample.meta.id = sample.meta.id
  771. img_metas = {}
  772. for key in self.meta_keys:
  773. if key in sample:
  774. img_metas[key] = sample[key]
  775. filtered_sample["img_metas"] = img_metas
  776. for key in self.keys:
  777. filtered_sample[key] = sample[key]
  778. return filtered_sample
  779. @benchmark.timeit
  780. class GetInferInput:
  781. """Collect infer input data from transformed sample"""
  782. def collate_fn(self, batch):
  783. sample = batch[0]
  784. collated_batch = {}
  785. collated_fields = [
  786. "img",
  787. "points",
  788. "img_metas",
  789. "gt_bboxes_3d",
  790. "gt_labels_3d",
  791. "modality",
  792. "meta",
  793. "idx",
  794. "img_depth",
  795. ]
  796. for k in list(sample.keys()):
  797. if k not in collated_fields:
  798. continue
  799. if k == "img":
  800. collated_batch[k] = np.stack([elem[k] for elem in batch], axis=0)
  801. elif k == "img_depth":
  802. collated_batch[k] = np.stack(
  803. [np.stack(elem[k], axis=0) for elem in batch], axis=0
  804. )
  805. else:
  806. collated_batch[k] = [elem[k] for elem in batch]
  807. return collated_batch
  808. def __call__(self, sample):
  809. """Call function to infer input data from transformed sample
  810. Args:
  811. sample (Sample): The input sample data.
  812. Returns:
  813. infer_input (list): A list containing all the input data for inference.
  814. sample_id (str): token id of the input sample.
  815. """
  816. if sample.modality == "multimodal" or sample.modality == "multiview":
  817. if "img" in sample.keys():
  818. sample.img = np.stack(
  819. [img.transpose(2, 0, 1) for img in sample.img], axis=0
  820. )
  821. sample = self.collate_fn([sample])
  822. infer_input = []
  823. img = sample.get("img", None)[0]
  824. infer_input.append(img.astype(np.float32))
  825. lidar2img = np.stack(sample["img_metas"][0]["lidar2img"]).astype(np.float32)
  826. infer_input.append(lidar2img)
  827. points = sample.get("points", None)[0]
  828. infer_input.append(points.astype(np.float32))
  829. img_metas = {
  830. "input_lidar_path": sample["img_metas"][0]["pts_filename"],
  831. "input_img_paths": sample["img_metas"][0]["filename"],
  832. "sample_id": sample["img_metas"][0]["sample_idx"],
  833. }
  834. return infer_input, img_metas