diff --git a/demo/estimate_smpl.py b/demo/estimate_smpl.py index d7ca6b2d..a60dd41b 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: @@ -238,8 +238,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:]) @@ -271,9 +270,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( @@ -412,8 +409,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]: @@ -463,9 +459,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 0a33e25f..deb9e7b0 100644 --- a/mmhuman3d/utils/demo_utils.py +++ b/mmhuman3d/utils/demo_utils.py @@ -11,6 +11,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 @@ -717,6 +718,28 @@ def get_different_colors(number_of_colors, 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) + + class RunningAverage(): r"""A helper class to calculate running average in a sliding window. 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: 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'))