Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
    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])