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])