def do_3d_tracking(self, detections_2d, detections_3d): start = time.time() #rospy.loginfo("Tracking frame") # convert_detections boxes_2d = [] boxes_3d = [] valid_3d = [] features_2d = [] features_3d = [] dets_2d = sorted(detections_2d.detection2d_with_features, key=lambda x: x.frame_det_id) dets_3d = sorted(detections_3d.detection3d_with_features, key=lambda x: x.frame_det_id) i, j = 0, 0 while i < len(dets_2d) and j < len(dets_3d): det_2d = dets_2d[i] det_3d = dets_3d[j] if det_2d.frame_det_id == det_3d.frame_det_id: i += 1 j += 1 valid_3d.append(det_3d.valid) boxes_2d.append( np.array([ det_2d.x1, det_2d.y1, det_2d.x2, det_2d.y2, 1, -1, -1 ])) features_2d.append(torch.Tensor(det_2d.feature).to('cuda:0')) if det_3d.valid: boxes_3d.append( np.array([ det_3d.x, det_3d.y, det_3d.z, det_3d.l, det_3d.h, det_3d.w, det_3d.theta ])) features_3d.append( torch.Tensor(det_3d.feature).to('cuda:0')) else: boxes_3d.append(None) features_3d.append(None) elif det_2d.frame_det_id < det_3d.frame_det_id: i += 1 else: j += 1 if not boxes_3d: boxes_3d = None features_3d, features_2d = combine_features( features_2d, features_3d, valid_3d, self.combination_model, depth_weight=self.depth_weight) detections = convert_detections(boxes_2d, features_3d, features_2d, boxes_3d) self.tracker.predict() self.tracker.update(None, detections) tracked_array = TrackedPersons() tracked_array.header.stamp = detections_3d.header.stamp tracked_array.header.frame_id = 'occam' for track in self.tracker.tracks: if not track.is_confirmed(): continue #print('Confirmed track!') pose_msg = Pose() tracked_person_msg = TrackedPerson() tracked_person_msg.header.stamp = detections_3d.header.stamp tracked_person_msg.header.frame_id = 'occam' tracked_person_msg.track_id = track.track_id if track.time_since_update < 2: tracked_person_msg.is_matched = True else: tracked_person_msg.is_matched = False bbox = track.to_tlwh3d() covariance = track.get_cov().reshape(-1).tolist() pose_msg.position.x = bbox[0] pose_msg.position.y = bbox[1] - bbox[4] / 2 pose_msg.position.z = bbox[2] pose_msg = PoseWithCovariance(pose=pose_msg, covariance=covariance) tracked_person_msg.pose = pose_msg tracked_array.tracks.append(tracked_person_msg) self.tracker_output_pub.publish(tracked_array)