def inference(
detection_cfg,
skeleton_cfg,
dataset_cfg,
gpus=1,
worker_per_gpu=1,
):
# get frame num
video_file = dataset_cfg.video_file
video_name = video_file.strip('/n').split('/')[-1]
video_frames = mmcv.VideoReader(video_file)
num_frames = len(video_frames)
del video_frames
data_cfg = skeleton_cfg.data_cfg
if data_cfg.save_video:
data_cfg.img_dir = os.path.join(data_cfg.save_dir,
'{}.img'.format(video_name))
if os.path.exists(data_cfg.img_dir):
import shutil
shutil.rmtree(data_cfg.img_dir)
os.makedirs(data_cfg.img_dir)
# cache model checkpoints
cache_checkpoint(detection_cfg.checkpoint_file)
cache_checkpoint(skeleton_cfg.checkpoint_file)
# multiprocess settings
context = mp.get_context('spawn')
result_queue = context.Queue(num_frames)
procs = []
for w in range(gpus * worker_per_gpu):
shred_list = list(range(w, num_frames, gpus * worker_per_gpu))
p = context.Process(target=worker,
args=(video_file, shred_list, detection_cfg,
skeleton_cfg, data_cfg, w % gpus,
result_queue))
p.start()
procs.append(p)
all_result = []
print('\nPose estimation start:')
prog_bar = ProgressBar(num_frames)
for i in range(num_frames):
t = result_queue.get()
all_result.append(t)
prog_bar.update()
for p in procs:
p.join()
if len(all_result) == num_frames and data_cfg.save_video:
print('\n\nGenerate video:')
video_path = os.path.join(data_cfg.save_dir, video_name)
mmcv.frames2video(data_cfg.img_dir,
video_path,
filename_tmpl='{:01d}.png')
print('Video was saved to {}'.format(video_path))
import IPython
IPython.embed()
def build(detection_cfg,
estimation_cfg,
tracker_cfg,
video_dir,
out_dir,
gpus=1,
worker_per_gpu=1,
video_max_length=10000,
category_annotation=None):
cache_checkpoint(detection_cfg.checkpoint_file)
cache_checkpoint(estimation_cfg.checkpoint_file)
if tracker_cfg is not None:
raise NotImplementedError
if not os.path.isdir(out_dir):
os.makedirs(out_dir)
if category_annotation is None:
video_categories = dict()
else:
with open(category_annotation) as f:
video_categories = json.load(f)['annotations']
inputs = Manager().Queue(video_max_length)
results = Manager().Queue(video_max_length)
num_worker = gpus * worker_per_gpu
procs = []
for i in range(num_worker):
p = Process(target=worker,
args=(inputs, results, i % gpus, detection_cfg,
estimation_cfg))
procs.append(p)
p.start()
video_file_list = os.listdir(video_dir)
prog_bar = ProgressBar(len(video_file_list))
for video_file in video_file_list:
#print('video_file:',video_file)
reader = mmcv.VideoReader(os.path.join(video_dir, video_file))
video_frames = reader[:video_max_length]
annotations = []
num_keypoints = -1
count_frame = 0
for i, image in enumerate(video_frames):
#print('input_i:',i)
if image is None:
continue
count_frame += 1
inputs.put((i, image))
#print('\ncount_frame:',count_frame)
#print('len_frame',len(video_frames))
for i in range(count_frame):
#print('frame:',i)
t = results.get()
#print('t:',t)
if not t['has_return']:
continue
num_person = len(t['joint_preds'])
assert len(t['person_bbox']) == num_person
for j in range(num_person):
keypoints = [[p[0], p[1], round(s[0], 2)] for p, s in zip(
t['joint_preds'][j].round().astype(int).tolist(),
t['joint_scores'][j].tolist())]
num_keypoints = len(keypoints)
person_info = dict(person_bbox=t['person_bbox']
[j].round().astype(int).tolist(),
frame_index=t['frame_index'],
id=j,
person_id=None,
keypoints=keypoints)
annotations.append(person_info)
# output results
annotations = sorted(annotations, key=lambda x: x['frame_index'])
category_id = video_categories[video_file][
'category_id'] if video_file in video_categories else -1
info = dict(video_name=video_file,
resolution=reader.resolution,
num_frame=len(video_frames),
num_keypoints=num_keypoints,
keypoint_channels=['x', 'y', 'score'],
version='1.0')
video_info = dict(info=info,
category_id=category_id,
annotations=annotations)
with open(os.path.join(out_dir, video_file + '.json'), 'w') as f:
json.dump(video_info, f)
prog_bar.update()
# send end signals
for p in procs:
inputs.put((-1, None))
# wait to finish
for p in procs:
p.join()
print('\nBuild skeleton dataset to {}.'.format(out_dir))
return video_info
def inference(detection_cfg,
estimation_cfg,
video_file,
gpus=1,
worker_per_gpu=1,
save_dir=None):
video_frames = mmcv.VideoReader(video_file)
all_result = []
print('\nPose estimation:')
# case for single process
if gpus == 1 and worker_per_gpu == 1:
model = init_pose_estimator(detection_cfg, estimation_cfg, device=0)
prog_bar = ProgressBar(len(video_frames))
for i, image in enumerate(video_frames):
res = inference_pose_estimator(model, image)
res['frame_index'] = i
if save_dir is not None:
res['render_image'] = render(image, res['joint_preds'],
res['person_bbox'],
detection_cfg.bbox_thre)
all_result.append(res)
prog_bar.update()
# case for multi-process
else:
cache_checkpoint(detection_cfg.checkpoint_file)
cache_checkpoint(estimation_cfg.checkpoint_file)
num_worker = gpus * worker_per_gpu
procs = []
inputs = Manager().Queue(len(video_frames))
results = Manager().Queue(len(video_frames))
for i, image in enumerate(video_frames):
inputs.put((i, image))
for i in range(num_worker):
p = Process(target=worker,
args=(inputs, results, i % gpus, detection_cfg,
estimation_cfg, save_dir is not None))
procs.append(p)
p.start()
for i in range(len(video_frames)):
t = results.get()
all_result.append(t)
if 'prog_bar' not in locals():
prog_bar = ProgressBar(len(video_frames))
prog_bar.update()
for p in procs:
p.join()
# sort results
all_result = sorted(all_result, key=lambda x: x['frame_index'])
# generate video
if (len(all_result) == len(video_frames)) and (save_dir is not None):
print('\n\nGenerate video:')
video_name = video_file.strip('/n').split('/')[-1]
video_path = os.path.join(save_dir, video_name)
vwriter = cv2.VideoWriter(video_path,
mmcv.video.io.VideoWriter_fourcc(*('mp4v')),
video_frames.fps, video_frames.resolution)
prog_bar = ProgressBar(len(video_frames))
for r in all_result:
vwriter.write(r['render_image'])
prog_bar.update()
vwriter.release()
print('\nVideo was saved to {}'.format(video_path))
return all_result
def build(inputs,
detection_cfg,
estimation_cfg,
tracker_cfg,
video_dir,
gpus=1,
video_max_length=10000,
category_annotation=None):
print('data build start')
cache_checkpoint(detection_cfg.checkpoint_file)
cache_checkpoint(estimation_cfg.checkpoint_file)
if category_annotation is None:
video_categories = dict()
else:
with open(category_annotation) as f:
video_categories = json.load(f)['annotations']
if tracker_cfg is not None:
raise NotImplementedError
pose_estimators = init_pose_estimator(detection_cfg,
estimation_cfg,
device=0)
video_file_list = []
get_all_file(video_dir, video_file_list)
prog_bar = ProgressBar(len(video_file_list))
for video_path in video_file_list:
video_file = os.path.basename(video_path)
reader = mmcv.VideoReader(video_path)
video_frames = reader[:video_max_length]
annotations = []
num_keypoints = -1
for i, image in enumerate(video_frames):
res = inference_pose_estimator(pose_estimators, image)
res['frame_index'] = i
if not res['has_return']:
continue
num_person = len(res['joint_preds'])
assert len(res['person_bbox']) == num_person
for j in range(num_person):
keypoints = [[p[0], p[1], round(s[0], 2)] for p, s in zip(
res['joint_preds'][j].round().astype(int).tolist(),
res['joint_scores'][j].tolist())]
num_keypoints = len(keypoints)
person_info = dict(person_bbox=res['person_bbox']
[j].round().astype(int).tolist(),
frame_index=res['frame_index'],
id=j,
person_id=None,
keypoints=keypoints)
annotations.append(person_info)
annotations = sorted(annotations, key=lambda x: x['frame_index'])
category_id = video_categories[video_file][
'category_id'] if video_file in video_categories else -1
info = dict(video_name=video_file,
resolution=reader.resolution,
num_frame=len(video_frames),
num_keypoints=num_keypoints,
keypoint_channels=['x', 'y', 'score'],
version='1.0')
video_info = dict(info=info,
category_id=category_id,
annotations=annotations)
inputs.put(video_info)
prog_bar.update()