Spaces:
Runtime error
Runtime error
File size: 3,204 Bytes
2d5fdd1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
# coding=utf-8
# Copyright 2020 The Google AI Perception Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Estimate AIST++ 3D keypoints."""
import os
import pickle
from absl import app
from absl import flags
from absl import logging
from aist_plusplus.loader import AISTDataset
import numpy as np
FLAGS = flags.FLAGS
flags.DEFINE_string(
'anno_dir',
'/usr/local/google/home/ruilongli/data/public/aist_plusplus_final/',
'input local dictionary for AIST++ annotations.')
flags.DEFINE_string(
'save_dir',
'/usr/local/google/home/ruilongli/data/public/aist_plusplus_final/keypoints3d/',
'output local dictionary that stores AIST++ 3D keypoints.')
np.random.seed(0)
def main(_):
aist_dataset = AISTDataset(anno_dir=FLAGS.anno_dir)
for seq_name, env_name in aist_dataset.mapping_seq2env.items():
logging.info('processing %s', seq_name)
# Load camera parameters
cgroup = AISTDataset.load_camera_group(aist_dataset.camera_dir, env_name)
# load 2D keypoints
keypoints2d, det_scores, _ = AISTDataset.load_keypoint2d(
aist_dataset.keypoint2d_dir, seq_name=seq_name)
nviews, nframes, _, _ = keypoints2d.shape
assert det_scores.shape[0] == nviews
assert det_scores.shape[1] == nframes
if seq_name == 'gBR_sBM_cAll_d04_mBR0_ch01':
keypoints2d[4] = np.nan # not synced view
if seq_name == 'gJB_sBM_cAll_d07_mJB3_ch05':
keypoints2d[6] = np.nan # size 640x480
# filter keypoints to select those best points
kpt_thre = 0.15
ignore_idxs = np.where(keypoints2d[:, :, :, 2] < kpt_thre)
keypoints2d[ignore_idxs[0], ignore_idxs[1], ignore_idxs[2], :] = np.nan
det_thre = 0.0
ignore_idxs = np.where(det_scores < det_thre)
keypoints2d[ignore_idxs[0], ignore_idxs[1], :, :] = np.nan
keypoints2d = keypoints2d[:, :, :, 0:2]
# 3D pose triangulation and temporal optimization.
bones = [
(5, 7), (7, 9), (6, 8), (8, 10), (11, 13), (13, 15), (12, 14),
(14, 16), (0, 1), (0, 2), (1, 2), (0, 3), (0, 4), (3, 4),
] # COCO-format bone constrains
keypoints3d = cgroup.triangulate(
keypoints2d.reshape(nviews, -1, 2)
).reshape(nframes, -1, 3)
keypoints3d_optim = cgroup.triangulate_optim(
keypoints2d, constraints=bones, verbose=True
).reshape(nframes, -1, 3)
# Save to pkl
os.makedirs(FLAGS.save_dir, exist_ok=True)
keypoints_file = os.path.join(FLAGS.save_dir, f'{seq_name}.pkl')
with open(keypoints_file, 'wb') as f:
pickle.dump({
'keypoints3d': keypoints3d,
'keypoints3d_optim': keypoints3d_optim,
}, f, protocol=pickle.HIGHEST_PROTOCOL)
if __name__ == '__main__':
app.run(main)
|