ego_to_world = ego_poses[frame.timestamp_micros].ego_pose
for box in boxes_3d:
box.transform_(ego_to_world, 'world')
all_measurements[box.obj_type].append(cvt_bbox3d_to_measurement(box))
# invoke tracking
for obj_type, manager in managers.items():
manager.run_(all_measurements[obj_type], frame_idx)
world_to_ego = np.linalg.inv(ego_to_world)
for obj_type, manager in managers.items():
for tracklet in manager.all_tracklets:
if tracklet.tail.stamp == frame_idx and not tracklet.just_born and tracklet.conf > 0.1:
box = cvt_state_to_bbox3d(tracklet.tail,
tracklet.id,
tracklet.most_recent_meas_score,
frame='world',
obj_type=obj_type)
# map back to ego_vehicle frame
box.transform_(world_to_ego, 'ego_vehicle')
# find camera this box is visible on and draw
if obj_type == 'VEHICLE':
find_camera(box, waymo_cameras)
draw_bbox3d_on_image(
box, waymo_cameras,
'{}:{}'.format(obj_type[:3], tracklet.id))
# ------------------------------------------------------------------
for cam in waymo_cameras:
_h = resized_height_front if 'FRONT' in cam.name else resized_height_back
cam.im = cv2.resize(cam.im, (resized_width, _h))
def generate_tracking_results_for_one_record(context_name, context_idx,
nbr_contexts):
"""Generate tracking results for one record represents by the context_name
Args:
context_name (str): name of the context
context_idx (int): context index
nbr_contexts (int): total number of contexts
Returns:
metrics_pb2.Objects: the object stores all the tracking result of this record
"""
print('\nGenerate tracking result for context {}/{}'.format(
context_idx, nbr_contexts))
# get path to ego_poses & detections file
ego_poses_file = os.path.join(data_root, context_name,
'{}_stamp_and_pose.txt'.format(context_name))
detections_file = os.path.join(
data_root, context_name,
'{}_unpacked_detection.txt'.format(context_name))
# parse these files to get the dict {timestamp_micros: information}
ego_poses = parse_ego_pose_file(ego_poses_file)
detections = parse_detection_file(detections_file)
# make sure timestamp are in ascending order
all_timestamps = sorted(list(ego_poses.keys()))
# init tracklet managers
managers = {
name: TrackletManager(name)
for name in waymo_to_nuscenes.keys()
}
tracking_results = metrics_pb2.Objects() # to store tracking results
# main loop
frame_idx = 0
for timestamp in tqdm(all_timestamps):
# get detection of this frame
boxes_3d = []
if timestamp in detections.keys():
boxes_3d = [
waymo_object_to_bbox3d(o, frame_idx)
for o in detections[timestamp]
]
# convert detection in 'ego_vehicle' frame to measurement in 'world' frame
ego_to_world = ego_poses[timestamp].ego_pose
all_measurements = {name: [] for name in managers.keys()}
for box in boxes_3d:
if box.score > 0.15:
box.transform_(ego_to_world, 'world')
all_measurements[box.obj_type].append(
cvt_bbox3d_to_measurement(box))
# invoke tracking
for obj_type, manager in managers.items():
manager.run_(all_measurements[obj_type], frame_idx)
# log tracking result
world_to_ego = np.linalg.inv(ego_to_world)
for obj_type, manager in managers.items():
for tracklet in manager.all_tracklets:
if tracklet.tail.stamp == frame_idx and not tracklet.just_born and tracklet.conf > 0.1:
box = cvt_state_to_bbox3d(tracklet.tail,
tracklet.id,
tracklet.most_recent_meas_score,
frame='world',
obj_type=obj_type)
box.transform_(
world_to_ego,
'ego_vehicle') # map back to ego_vehicle frame
o = format_tracking_result(box, context_name, timestamp)
tracking_results.objects.append(o)
# move on
frame_idx += 1
print('writing tracking results of {}'.format(context_name))
with open(
os.path.join(result_root,
'{}_tracking_result.bin'.format(context_name)),
'wb') as f:
f.write(tracking_results.SerializeToString())
print(
'----------------------------------------------------------------------------\n'
)
def main(seq_name, tracked_obj_class, create_video=True):
print("Demo tracking {} on sequence {}. Press 'q' to stop".format(
tracked_obj_class, seq_name))
# get OXTS data
oxts_file = os.path.join(kitti_root, 'data_tracking_oxts', 'training',
'oxts', '{}.txt'.format(seq_name))
oxts = load_oxts_packets_and_poses([oxts_file])
# get sequence image
img_dir = os.path.join(kitti_root, 'data_tracking_image_2', 'training',
'image_02', seq_name)
img_names = [file for file in os.listdir(img_dir) if file.endswith('png')]
img_names.sort()
if create_video:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
img = cv2.imread(os.path.join(img_dir, img_names[0]))
video = cv2.VideoWriter('video.avi', fourcc, 1,
(img.shape[1], img.shape[0]))
# get sensor calib
calib_file = os.path.join(kitti_root, 'data_tracking_calib', 'training',
'calib', '{}.txt'.format(seq_name))
calib = Calibration(calib_file)
# get detection
detect_file = './data/kitti/pointrcnn_detection/pointrcnn_{}_val/{}.txt'.format(
tracked_obj_class, seq_name)
detection = parse_detection_file(detect_file)
# init tracklet manager
manager = TrackletManager(tracked_obj_class)
# main loop
world_to_c0 = np.eye(4)
for frame_idx, name in enumerate(img_names):
img = cv2.imread(os.path.join(img_dir, name))
# get ego vehicle pose
imu_to_world = oxts[frame_idx].imu_to_world
if frame_idx == 0:
# get mapping from world to c0
c0_to_world = imu_to_world @ calib.cam_to_imu
world_to_c0 = np.linalg.inv(c0_to_world)
ci_to_world = imu_to_world @ calib.cam_to_imu
ci_to_c0 = world_to_c0 @ ci_to_world
# get 3D object detection in camera frame
boxes_3d = []
if frame_idx in detection.keys():
boxes_3d = [
kitti_obj_to_bbox3d(o, timestamp=frame_idx, obj_type='Car')
for o in detection[frame_idx]
]
# map boxes from current camera frame to first camera frame
for box in boxes_3d:
box.transform_(ci_to_c0, 'c0')
# ------------------------------------------------------------------ #
# tracking code
# should go
# here
# convert boxes to measurements
all_measurements = [cvt_bbox3d_to_measurement(box) for box in boxes_3d]
manager.run_(all_measurements, frame_idx)
# ------------------------------------------------------------------ #
c0_to_cam = np.linalg.inv(ci_to_c0)
for tracklet in manager.all_tracklets:
if tracklet.tail.stamp == frame_idx and not tracklet.just_born:
# convert tracklet's tail to a Bbox3D
box = cvt_state_to_bbox3d(tracklet.tail, tracklet.id,
tracklet.most_recent_meas_score)
# map box back to camera frame for drawing
box.transform_(c0_to_cam, 'camera')
# draw boxes
box_to_cam = get_box_to_cam_trans(box)
proj = box.project_on_image(box_to_cam, calib.cam_proj_mat)
draw_bbox3d(img, proj, tracklet.id)
cv2.imshow('left_color', img)
if create_video:
video.write(img)
if cv2.waitKey(100) & 0xFF == ord('q'):
break
cv2.destroyAllWindows()
if create_video:
video.release()
def main(scene_index):
print(
"Demo tracking on scene {} of NuScenes mini. Press 'q' to stop".format(
scene_index))
nusc = NuScenes(dataroot=nuscenes_root_mini,
version='v1.0-mini',
verbose=False)
detection_dir = './data/nuscenes/megvii_detection'
scene = nusc.scene[scene_index]
print('scene token: ', scene['token'])
detection_file = os.path.join(detection_dir,
'{}.txt'.format(scene['token']))
detections = parse_detection_file(detection_file)
# init tracklet managers
tracklet_managers = {
obj_type: TrackletManager(obj_type)
for obj_type in nuscenes_tracking_names
}
# prepare for rendering 6 cameras
imsize = (640, 360)
window_name = '{}'.format(scene['name'])
cv2.namedWindow(window_name)
cv2.moveWindow(window_name, 0, 0)
canvas = np.ones((2 * imsize[1], 3 * imsize[0], 3), np.uint8)
layout = {
'CAM_FRONT_LEFT': (0, 0),
'CAM_FRONT': (imsize[0], 0),
'CAM_FRONT_RIGHT': (2 * imsize[0], 0),
'CAM_BACK_LEFT': (0, imsize[1]),
'CAM_BACK': (imsize[0], imsize[1]),
'CAM_BACK_RIGHT': (2 * imsize[0], imsize[1]),
}
horizontal_flip = ['CAM_BACK_LEFT', 'CAM_BACK',
'CAM_BACK_RIGHT'] # Flip these for aesthetic reasons.
frame_idx = 0
sample_token = scene['first_sample_token']
while sample_token != '':
sample = nusc.get('sample', sample_token)
# get camera information
camera_tokens = [
token for k, token in sample['data'].items() if 'CAM' in k
]
cam_infos = tuple(
[NuCamera(cam_token, nusc) for cam_token in camera_tokens])
# get detection of this frame
boxes_3d = []
if frame_idx in detections.keys():
boxes_3d = [
nuscenes_object_to_bbox3d(o) for o in detections[frame_idx]
]
# ------------------------------------------------------------------#
# tracking code
# should go
# here
# convert boxes to measurements
all_measurements = {
obj_type: []
for obj_type in nuscenes_tracking_names
}
for box in boxes_3d:
all_measurements[box.obj_type].append(
cvt_bbox3d_to_measurement(box))
for obj_type, manager in tracklet_managers.items():
manager.run_(all_measurements[obj_type], frame_idx)
# ------------------------------------------------------------------#
# draw boxes on camera images
for obj_type, manager in tracklet_managers.items():
for tracklet in manager.all_tracklets:
if tracklet.tail.stamp == frame_idx and not tracklet.just_born:
box = cvt_state_to_bbox3d(tracklet.tail,
tracklet.id,
tracklet.most_recent_meas_score,
frame='world')
# find camera this box is visible on and draw
find_camera(box, cam_infos)
draw_bbox3d_on_image(
box, cam_infos, '{}:{}'.format(obj_type[:3],
tracklet.id))
for cam in cam_infos:
cam.im = cv2.resize(cam.im, imsize)
if cam.name in horizontal_flip:
cam.im = cam.im[:, ::-1, :]
canvas[layout[cam.name][1]:layout[cam.name][1] + imsize[1],
layout[cam.name][0]:layout[cam.name][0] +
imsize[0], :] = cam.im
cv2.imshow(window_name, canvas)
if cv2.waitKey(300) & 0xFF == ord('q'):
break
# move on
frame_idx += 1
sample_token = sample['next']
cv2.destroyAllWindows()
for box in boxes_3d:
box.transform_(ci_to_c0, 'c0')
# invoke tracking functionalities
all_measurements = [
cvt_bbox3d_to_measurement(box) for box in boxes_3d
]
manager.run_(all_measurements, frame_idx)
# log result
c0_to_cam = np.linalg.inv(ci_to_c0)
for tracklet in manager.all_tracklets:
if tracklet.tail.stamp == frame_idx and not tracklet.just_born:
# convert tracklet's tail to a Bbox3D
box = cvt_state_to_bbox3d(tracklet.tail, tracklet.id,
tracklet.most_recent_meas_score,
tracklet.kitti_meas_alpha)
# map box back to camera frame for drawing
box.transform_(c0_to_cam, 'camera')
# project box onto image for 2D tracking benchmarking
box_to_cam = get_box_to_cam_trans(box)
proj = box.project_on_image(box_to_cam, calib.cam_proj_mat)
# write result
box_2d = np.append(np.amin(proj, axis=0),
np.amax(proj, axis=0)).tolist()
result_file.write(
format_result(frame_idx, tracklet.id, obj_type, box,
box_2d))
result_file.close()
print('\n--------------------------------------------')