From 346e7f0ce6897a1a4d529ef917c7a3809512ac08 Mon Sep 17 00:00:00 2001 From: rlleshi Date: Thu, 14 Jul 2022 12:47:30 +0200 Subject: [PATCH 1/2] add batch processing --- demo/estimate_smpl.py | 16 +++++----------- mmhuman3d/utils/demo_utils.py | 23 +++++++++++++++++++++++ mmhuman3d/utils/ffmpeg_utils.py | 2 +- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/demo/estimate_smpl.py b/demo/estimate_smpl.py index 8e51fc9b..25a3e587 100644 --- a/demo/estimate_smpl.py +++ b/demo/estimate_smpl.py @@ -18,6 +18,7 @@ from mmhuman3d.core.visualization.visualize_smpl import visualize_smpl_hmr from mmhuman3d.data.data_structures.human_data import HumanData from mmhuman3d.utils.demo_utils import ( + batch_array_to_images, extract_feature_sequence, get_speed_up_interval, prepare_frames, @@ -27,7 +28,6 @@ speed_up_interpolate, speed_up_process, ) -from mmhuman3d.utils.ffmpeg_utils import array_to_images from mmhuman3d.utils.transforms import rotmat_to_aa try: @@ -237,8 +237,7 @@ def single_person_with_mmdet(args, frames_iter): human_data = HumanData() frames_folder = osp.join(args.output, 'images') os.makedirs(frames_folder, exist_ok=True) - array_to_images( - np.array(frames_iter)[frame_id_list], output_folder=frames_folder) + batch_array_to_images(frames_iter, frame_id_list, frames_folder) for i, img_i in enumerate(sorted(os.listdir(frames_folder))): body_pose_.append(smpl_poses[i][1:]) @@ -270,9 +269,7 @@ def single_person_with_mmdet(args, frames_iter): else: frames_folder = osp.join(Path(args.show_path).parent, 'images') os.makedirs(frames_folder, exist_ok=True) - array_to_images( - np.array(frames_iter)[frame_id_list], - output_folder=frames_folder) + batch_array_to_images(frames_iter, frame_id_list, frames_folder) body_model_config = dict(model_path=args.body_model_dir, type='smpl') visualize_smpl_hmr( @@ -410,8 +407,7 @@ def multi_person_with_mmtracking(args, frames_iter): human_data = HumanData() frames_folder = osp.join(args.output, 'images') os.makedirs(frames_folder, exist_ok=True) - array_to_images( - np.array(frames_iter)[frame_id_list], output_folder=frames_folder) + batch_array_to_images(frames_iter, frame_id_list, frames_folder) for i, img_i in enumerate(sorted(os.listdir(frames_folder))): for person_i in track_ids_lists[i]: @@ -461,9 +457,7 @@ def multi_person_with_mmtracking(args, frames_iter): else: frames_folder = osp.join(Path(args.show_path).parent, 'images') os.makedirs(frames_folder, exist_ok=True) - array_to_images( - np.array(frames_iter)[frame_id_list], - output_folder=frames_folder) + batch_array_to_images(frames_iter, frame_id_list, frames_folder) body_model_config = dict(model_path=args.body_model_dir, type='smpl') visualize_smpl_hmr( poses=compressed_poses.reshape(-1, max_instance, 24 * 3), diff --git a/mmhuman3d/utils/demo_utils.py b/mmhuman3d/utils/demo_utils.py index 8fc03aec..d7cb9f03 100644 --- a/mmhuman3d/utils/demo_utils.py +++ b/mmhuman3d/utils/demo_utils.py @@ -7,6 +7,7 @@ from scipy import interpolate from mmhuman3d.core.post_processing import build_post_processing +from mmhuman3d.utils.ffmpeg_utils import array_to_images try: from typing import Literal @@ -707,3 +708,25 @@ def get_different_colors(number_of_colors, colors_final.append(color_dict[channel]) colors_final = np.concatenate(colors_final, -1) return colors_final + + +def batch_array_to_images(frames_iter: list, + frame_id_list: list, + frames_folder: str, + splits: int = 10): + """Batch process the total amount of frames to avoid oom problems that + might arise when loading a large number of images to memory. + + Args: + frames_iter (list): image frames list + frame_id_list (list): image frames indices + frames_folder (str): frames folder + splits (int): number of array splits + """ + + def frames_array_split(frames_iter: list, frame_id_list: list): + for f_i_l in np.array_split(frame_id_list, splits): + yield np.array(frames_iter)[f_i_l] + + for frames_iter_split in frames_array_split(frames_iter, frame_id_list): + array_to_images(frames_iter_split, output_folder=frames_folder) diff --git a/mmhuman3d/utils/ffmpeg_utils.py b/mmhuman3d/utils/ffmpeg_utils.py index 0c30b5db..d00b51bf 100644 --- a/mmhuman3d/utils/ffmpeg_utils.py +++ b/mmhuman3d/utils/ffmpeg_utils.py @@ -230,7 +230,7 @@ def array_to_images( '-f', 'image2', '-start_number', - '0', + str(len(os.listdir(output_folder))), os.path.join(output_folder, img_format), ] if not disable_log: From eb4285ee49cfa5c023d5d3e10d73e9b239d48b48 Mon Sep 17 00:00:00 2001 From: rlleshi Date: Tue, 9 Aug 2022 23:43:22 +0200 Subject: [PATCH 2/2] add ut --- tests/test_vis/test_ffmpeg.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_vis/test_ffmpeg.py b/tests/test_vis/test_ffmpeg.py index 8188e48c..28607115 100644 --- a/tests/test_vis/test_ffmpeg.py +++ b/tests/test_vis/test_ffmpeg.py @@ -4,6 +4,7 @@ import numpy as np import pytest +from mmhuman3d.utils.demo_utils import batch_array_to_images from mmhuman3d.utils.ffmpeg_utils import ( array_to_images, array_to_video, @@ -98,7 +99,8 @@ def test_array_saver(): v = np.random.randint( low=0, high=255, size=(30, 512, 512, 3), dtype=np.uint8) - array_to_images(v, osp.join(root, 'img_folder_saver')) + batch_array_to_images(v, list(range(len(v))), + osp.join(root, 'img_folder_saver')) array_to_video(v, osp.join(root, 'test_saver.mp4')) array_to_video(v[:, :-1, :-1, :], osp.join(root, 'test_even.mp4')) assert os.path.isfile(osp.join(root, 'test_saver.mp4'))