def data_frame_to_object_annotations(annotations, model, has_z=False, frame_id=0, resolution=None): is_valid_model(model) if annotations.empty: return ObjectAnnotations() n_model_joints = int(model.strip('joints')) joints_values = len( annotations.drop(['sample_id', 'person_id'], axis=1).columns) n_annotations_joints = int(joints_values / (4 if has_z else 3)) n_joints = min(n_model_joints, n_annotations_joints) annotations_pb = ObjectAnnotations() annotations_pb.frame_id = frame_id if resolution is not None: annotations_pb.resolution.width = resolution[0] annotations_pb.resolution.height = resolution[1] for _, annotation in annotations.iterrows(): skeleton = annotations_pb.objects.add() skeleton.id = int(annotation['person_id']) for joint_id in range(n_joints): x = annotation['j{:d}x'.format(joint_id)] y = annotation['j{:d}y'.format(joint_id)] z = annotation['j{:d}z'.format(joint_id)] if has_z else 0.0 c = annotation['j{:d}c'.format(joint_id)] human_keypoint = index_to_human_keypoint(joint_id, model) # check for invalid joint if (x == 0.0 and y == 0.0 and z == 0.0) or c < 0.0: continue if human_keypoint == HKP.Value('UNKNOWN_HUMAN_KEYPOINT'): continue keypoint = skeleton.keypoints.add() keypoint.position.x = x keypoint.position.y = y keypoint.position.z = z keypoint.score = c keypoint.id = human_keypoint return annotations_pb
detections_ids.sort() output_file = os.path.join(dataset_folder, '{}_2d_dataset'.format(model)) wirter_2D = ProtobufWriter(output_file) output_file = os.path.join(dataset_folder, '{}_3d_dataset'.format(model)) wirter_3D = ProtobufWriter(output_file) log.info("[{}] {} sequences", dataset, len(detections_ids)) for detection_id in detections_ids: detections_file = os.path.join(detections_folder, DETECTION_FILE.format(detection_id)) with open(detections_file, 'r') as f: detections = json.load(f) sks_3D = ObjectAnnotations() sks_3D.frame_id = referencial sks_2D = {} for camera in cameras: sks_2D[camera] = ObjectAnnotations() sks_2D[camera].resolution.width = calibs[camera].resolution.width sks_2D[camera].resolution.height = calibs[camera].resolution.height for body in detections['bodies']: joints_3D = np.array(body[joints_key]).reshape( (-1, 4)).transpose()[0:3, :] sk_3D, invalid_joints = joints_to_object_annotation( joints_3D, model) sk_3D.id = body[ 'id'] # on sk_3D, ObjectAnnotation::id stores person id from dataset sks_3D.objects.extend([sk_3D])