You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

kitti_data_utils.py 12 kB

2 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324
  1. """
  2. # -*- coding: utf-8 -*-
  3. -----------------------------------------------------------------------------------
  4. # Author: Nguyen Mau Dung
  5. # DoC: 2020.08.17
  6. # email: nguyenmaudung93.kstn@gmail.com
  7. -----------------------------------------------------------------------------------
  8. # Description: The utils of the kitti dataset
  9. """
  10. from __future__ import print_function
  11. import os
  12. import sys
  13. import numpy as np
  14. import cv2
  15. src_dir = os.path.dirname(os.path.realpath(__file__))
  16. # while not src_dir.endswith("sfa"):
  17. # src_dir = os.path.dirname(src_dir)
  18. if src_dir not in sys.path:
  19. sys.path.append(src_dir)
  20. import config.kitti_config as cnf
  21. class Object3d(object):
  22. ''' 3d object label '''
  23. def __init__(self, label_file_line):
  24. data = label_file_line.split(' ')
  25. data[1:] = [float(x) for x in data[1:]]
  26. # extract label, truncation, occlusion
  27. self.type = data[0] # 'Car', 'Pedestrian', ...
  28. self.cls_id = self.cls_type_to_id(self.type)
  29. self.truncation = data[1] # truncated pixel ratio [0..1]
  30. self.occlusion = int(data[2]) # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown
  31. self.alpha = data[3] # object observation angle [-pi..pi]
  32. # extract 2d bounding box in 0-based coordinates
  33. self.xmin = data[4] # left
  34. self.ymin = data[5] # top
  35. self.xmax = data[6] # right
  36. self.ymax = data[7] # bottom
  37. self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])
  38. # extract 3d bounding box information
  39. self.h = data[8] # box height
  40. self.w = data[9] # box width
  41. self.l = data[10] # box length (in meters)
  42. self.t = (data[11], data[12], data[13]) # location (x,y,z) in camera coord.
  43. self.dis_to_cam = np.linalg.norm(self.t)
  44. self.ry = data[14] # yaw angle (around Y-axis in camera coordinates) [-pi..pi]
  45. self.score = data[15] if data.__len__() == 16 else -1.0
  46. self.level_str = None
  47. self.level = self.get_obj_level()
  48. def cls_type_to_id(self, cls_type):
  49. if cls_type not in cnf.CLASS_NAME_TO_ID.keys():
  50. return -1
  51. return cnf.CLASS_NAME_TO_ID[cls_type]
  52. def get_obj_level(self):
  53. height = float(self.box2d[3]) - float(self.box2d[1]) + 1
  54. if height >= 40 and self.truncation <= 0.15 and self.occlusion <= 0:
  55. self.level_str = 'Easy'
  56. return 1 # Easy
  57. elif height >= 25 and self.truncation <= 0.3 and self.occlusion <= 1:
  58. self.level_str = 'Moderate'
  59. return 2 # Moderate
  60. elif height >= 25 and self.truncation <= 0.5 and self.occlusion <= 2:
  61. self.level_str = 'Hard'
  62. return 3 # Hard
  63. else:
  64. self.level_str = 'UnKnown'
  65. return 4
  66. def print_object(self):
  67. print('Type, truncation, occlusion, alpha: %s, %d, %d, %f' % \
  68. (self.type, self.truncation, self.occlusion, self.alpha))
  69. print('2d bbox (x0,y0,x1,y1): %f, %f, %f, %f' % \
  70. (self.xmin, self.ymin, self.xmax, self.ymax))
  71. print('3d bbox h,w,l: %f, %f, %f' % \
  72. (self.h, self.w, self.l))
  73. print('3d bbox location, ry: (%f, %f, %f), %f' % \
  74. (self.t[0], self.t[1], self.t[2], self.ry))
  75. def to_kitti_format(self):
  76. kitti_str = '%s %.2f %d %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f' \
  77. % (self.type, self.truncation, int(self.occlusion), self.alpha, self.box2d[0], self.box2d[1],
  78. self.box2d[2], self.box2d[3], self.h, self.w, self.l, self.t[0], self.t[1], self.t[2],
  79. self.ry, self.score)
  80. return kitti_str
  81. def read_label(label_filename):
  82. lines = [line.rstrip() for line in open(label_filename)]
  83. objects = [Object3d(line) for line in lines]
  84. return objects
  85. class Calibration(object):
  86. ''' Calibration matrices and utils
  87. 3d XYZ in <label>.txt are in rect camera coord.
  88. 2d box xy are in image2 coord
  89. Points in <lidar>.bin are in Velodyne coord.
  90. y_image2 = P^2_rect * x_rect
  91. y_image2 = P^2_rect * R0_rect * Tr_velo_to_cam * x_velo
  92. x_ref = Tr_velo_to_cam * x_velo
  93. x_rect = R0_rect * x_ref
  94. P^2_rect = [f^2_u, 0, c^2_u, -f^2_u b^2_x;
  95. 0, f^2_v, c^2_v, -f^2_v b^2_y;
  96. 0, 0, 1, 0]
  97. = K * [1|t]
  98. image2 coord:
  99. ----> x-axis (u)
  100. |
  101. |
  102. v y-axis (v)
  103. velodyne coord:
  104. front x, left y, up z
  105. rect/ref camera coord:
  106. right x, down y, front z
  107. Ref (KITTI paper): http://www.cvlibs.net/publications/Geiger2013IJRR.pdf
  108. TODO(rqi): do matrix multiplication only once for each projection.
  109. '''
  110. def __init__(self, calib_filepath):
  111. calibs = self.read_calib_file(calib_filepath)
  112. # Projection matrix from rect camera coord to image2 coord
  113. self.P2 = calibs['P2']
  114. self.P2 = np.reshape(self.P2, [3, 4])
  115. self.P3 = calibs['P3']
  116. self.P3 = np.reshape(self.P3, [3, 4])
  117. # Rigid transform from Velodyne coord to reference camera coord
  118. self.V2C = calibs['Tr_velo2cam']
  119. self.V2C = np.reshape(self.V2C, [3, 4])
  120. # Rotation from reference camera coord to rect camera coord
  121. self.R0 = calibs['R_rect']
  122. self.R0 = np.reshape(self.R0, [3, 3])
  123. # Camera intrinsics and extrinsics
  124. self.c_u = self.P2[0, 2]
  125. self.c_v = self.P2[1, 2]
  126. self.f_u = self.P2[0, 0]
  127. self.f_v = self.P2[1, 1]
  128. self.b_x = self.P2[0, 3] / (-self.f_u) # relative
  129. self.b_y = self.P2[1, 3] / (-self.f_v)
  130. def read_calib_file(self, filepath):
  131. with open(filepath) as f:
  132. lines = f.readlines()
  133. obj = lines[2].strip().split(' ')[1:]
  134. P2 = np.array(obj, dtype=np.float32)
  135. obj = lines[3].strip().split(' ')[1:]
  136. P3 = np.array(obj, dtype=np.float32)
  137. obj = lines[4].strip().split(' ')[1:]
  138. R0 = np.array(obj, dtype=np.float32)
  139. obj = lines[5].strip().split(' ')[1:]
  140. Tr_velo_to_cam = np.array(obj, dtype=np.float32)
  141. return {'P2': P2.reshape(3, 4),
  142. 'P3': P3.reshape(3, 4),
  143. 'R_rect': R0.reshape(3, 3),
  144. 'Tr_velo2cam': Tr_velo_to_cam.reshape(3, 4)}
  145. def cart2hom(self, pts_3d):
  146. """
  147. :param pts: (N, 3 or 2)
  148. :return pts_hom: (N, 4 or 3)
  149. """
  150. pts_hom = np.hstack((pts_3d, np.ones((pts_3d.shape[0], 1), dtype=np.float32)))
  151. return pts_hom
  152. def compute_radius(det_size, min_overlap=0.7):
  153. height, width = det_size
  154. a1 = 1
  155. b1 = (height + width)
  156. c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
  157. sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
  158. r1 = (b1 + sq1) / 2
  159. a2 = 4
  160. b2 = 2 * (height + width)
  161. c2 = (1 - min_overlap) * width * height
  162. sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
  163. r2 = (b2 + sq2) / 2
  164. a3 = 4 * min_overlap
  165. b3 = -2 * min_overlap * (height + width)
  166. c3 = (min_overlap - 1) * width * height
  167. sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
  168. r3 = (b3 + sq3) / 2
  169. return min(r1, r2, r3)
  170. def gaussian2D(shape, sigma=1):
  171. m, n = [(ss - 1.) / 2. for ss in shape]
  172. y, x = np.ogrid[-m:m + 1, -n:n + 1]
  173. h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
  174. h[h < np.finfo(h.dtype).eps * h.max()] = 0
  175. return h
  176. def gen_hm_radius(heatmap, center, radius, k=1):
  177. diameter = 2 * radius + 1
  178. gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
  179. x, y = int(center[0]), int(center[1])
  180. height, width = heatmap.shape[0:2]
  181. left, right = min(x, radius), min(width - x, radius + 1)
  182. top, bottom = min(y, radius), min(height - y, radius + 1)
  183. masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right]
  184. masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
  185. if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
  186. np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
  187. return heatmap
  188. def get_filtered_lidar(lidar, boundary, labels=None):
  189. minX = boundary['minX']
  190. maxX = boundary['maxX']
  191. minY = boundary['minY']
  192. maxY = boundary['maxY']
  193. minZ = boundary['minZ']
  194. maxZ = boundary['maxZ']
  195. # Remove the point out of range x,y,z
  196. mask = np.where((lidar[:, 0] >= minX) & (lidar[:, 0] <= maxX) &
  197. (lidar[:, 1] >= minY) & (lidar[:, 1] <= maxY) &
  198. (lidar[:, 2] >= minZ) & (lidar[:, 2] <= maxZ))
  199. lidar = lidar[mask]
  200. lidar[:, 2] = lidar[:, 2] - minZ
  201. if labels is not None:
  202. label_x = (labels[:, 1] >= minX) & (labels[:, 1] < maxX)
  203. label_y = (labels[:, 2] >= minY) & (labels[:, 2] < maxY)
  204. label_z = (labels[:, 3] >= minZ) & (labels[:, 3] < maxZ)
  205. mask_label = label_x & label_y & label_z
  206. labels = labels[mask_label]
  207. return lidar, labels
  208. else:
  209. return lidar
  210. def box3d_corners_to_center(box3d_corner):
  211. # (N, 8, 3) -> (N, 7)
  212. assert box3d_corner.ndim == 3
  213. xyz = np.mean(box3d_corner, axis=1)
  214. h = abs(np.mean(box3d_corner[:, 4:, 2] - box3d_corner[:, :4, 2], axis=1, keepdims=True))
  215. w = (np.sqrt(np.sum((box3d_corner[:, 0, [0, 1]] - box3d_corner[:, 1, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  216. np.sqrt(np.sum((box3d_corner[:, 2, [0, 1]] - box3d_corner[:, 3, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  217. np.sqrt(np.sum((box3d_corner[:, 4, [0, 1]] - box3d_corner[:, 5, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  218. np.sqrt(np.sum((box3d_corner[:, 6, [0, 1]] - box3d_corner[:, 7, [0, 1]]) ** 2, axis=1, keepdims=True))) / 4
  219. l = (np.sqrt(np.sum((box3d_corner[:, 0, [0, 1]] - box3d_corner[:, 3, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  220. np.sqrt(np.sum((box3d_corner[:, 1, [0, 1]] - box3d_corner[:, 2, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  221. np.sqrt(np.sum((box3d_corner[:, 4, [0, 1]] - box3d_corner[:, 7, [0, 1]]) ** 2, axis=1, keepdims=True)) +
  222. np.sqrt(np.sum((box3d_corner[:, 5, [0, 1]] - box3d_corner[:, 6, [0, 1]]) ** 2, axis=1, keepdims=True))) / 4
  223. yaw = (np.arctan2(box3d_corner[:, 2, 1] - box3d_corner[:, 1, 1],
  224. box3d_corner[:, 2, 0] - box3d_corner[:, 1, 0]) +
  225. np.arctan2(box3d_corner[:, 3, 1] - box3d_corner[:, 0, 1],
  226. box3d_corner[:, 3, 0] - box3d_corner[:, 0, 0]) +
  227. np.arctan2(box3d_corner[:, 2, 0] - box3d_corner[:, 3, 0],
  228. box3d_corner[:, 3, 1] - box3d_corner[:, 2, 1]) +
  229. np.arctan2(box3d_corner[:, 1, 0] - box3d_corner[:, 0, 0],
  230. box3d_corner[:, 0, 1] - box3d_corner[:, 1, 1]))[:, np.newaxis] / 4
  231. return np.concatenate([h, w, l, xyz, yaw], axis=1).reshape(-1, 7)
  232. def box3d_center_to_conners(box3d_center):
  233. h, w, l, x, y, z, yaw = box3d_center
  234. Box = np.array([[-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2],
  235. [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2],
  236. [0, 0, 0, 0, h, h, h, h]])
  237. rotMat = np.array([
  238. [np.cos(yaw), -np.sin(yaw), 0.0],
  239. [np.sin(yaw), np.cos(yaw), 0.0],
  240. [0.0, 0.0, 1.0]])
  241. velo_box = np.dot(rotMat, Box)
  242. cornerPosInVelo = velo_box + np.tile(np.array([x, y, z]), (8, 1)).T
  243. box3d_corner = cornerPosInVelo.transpose()
  244. return box3d_corner.astype(np.float32)
  245. if __name__ == '__main__':
  246. heatmap = np.zeros((96, 320))
  247. h, w = 40, 50
  248. radius = compute_radius((h, w))
  249. radius = max(0, int(radius))
  250. print('h: {}, w: {}, radius: {}, sigma: {}'.format(h, w, radius, (2 * radius + 1) / 6.))
  251. gen_hm_radius(heatmap, center=(200, 50), radius=radius)
  252. while True:
  253. cv2.imshow('heatmap', heatmap)
  254. if cv2.waitKey(0) & 0xff == 27:
  255. break
  256. max_pos = np.unravel_index(heatmap.argmax(), shape=heatmap.shape)
  257. print('max_pos: {}'.format(max_pos))

一站式算法开发平台、高性能分布式深度学习框架、先进算法模型库、视觉模型炼知平台、数据可视化分析平台等一系列平台及工具,在模型高效分布式训练、数据处理和可视分析、模型炼知和轻量化等技术上形成独特优势,目前已在产学研等各领域近千家单位及个人提供AI应用赋能