data_preprocess.py

# pylint: disable=W1203,W0718
"""
This module is used to process videos to prepare data for training. It utilizes various libraries and models
to perform tasks such as video frame extraction, audio extraction, face mask generation, and face embedding extraction.
The script takes in command-line arguments to specify the input and output directories, GPU status, level of parallelism,
and rank for distributed processing.

Usage:
 python -m scripts.data_preprocess --input_dir /path/to/video_dir --dataset_name dataset_name --gpu_status --parallelism 4 --rank 0

Example:
 python -m scripts.data_preprocess -i data/videos -o data/output -g -p 4 -r 0
"""
importargparse
importlogging
importos
frompathlibimportPath
fromtypingimportList

importcv2
importtorch
fromtqdmimporttqdm

fromhallo.datasets.audio_processorimportAudioProcessor
fromhallo.datasets.image_processorimportImageProcessorForDataProcessing
fromhallo.utils.utilimportconvert_video_to_images, extract_audio_from_videos

# Configure logging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s')


defsetup_directories(video_path: Path) ->dict:
"""
 Setup directories for storing processed files.

 Args:
 video_path (Path): Path to the video file.

 Returns:
 dict: A dictionary containing paths for various directories.
 """
base_dir=video_path.parent.parent
dirs= {
"face_mask": base_dir/"face_mask",
"sep_pose_mask": base_dir/"sep_pose_mask",
"sep_face_mask": base_dir/"sep_face_mask",
"sep_lip_mask": base_dir/"sep_lip_mask",
"face_emb": base_dir/"face_emb",
"audio_emb": base_dir/"audio_emb"
 }

forpathindirs.values():
path.mkdir(parents=True, exist_ok=True)

returndirs


defprocess_single_video(video_path: Path,
output_dir: Path,
image_processor: ImageProcessorForDataProcessing,
audio_processor: AudioProcessor,
step: int) ->None:
"""
 Process a single video file.

 Args:
 video_path (Path): Path to the video file.
 output_dir (Path): Directory to save the output.
 image_processor (ImageProcessorForDataProcessing): Image processor object.
 audio_processor (AudioProcessor): Audio processor object.
 gpu_status (bool): Whether to use GPU for processing.
 """
assertvideo_path.exists(), f"Video path {video_path} does not exist"
dirs=setup_directories(video_path)
logging.info(f"Processing video: {video_path}")

try:
ifstep==1:
images_output_dir=output_dir/'images'/video_path.stem
images_output_dir.mkdir(parents=True, exist_ok=True)
images_output_dir=convert_video_to_images(
video_path, images_output_dir)
logging.info(f"Images saved to: {images_output_dir}")

audio_output_dir=output_dir/'audios'
audio_output_dir.mkdir(parents=True, exist_ok=True)
audio_output_path=audio_output_dir/f'{video_path.stem}.wav'
audio_output_path=extract_audio_from_videos(
video_path, audio_output_path)
logging.info(f"Audio extracted to: {audio_output_path}")

face_mask, _, sep_pose_mask, sep_face_mask, sep_lip_mask=image_processor.preprocess(
images_output_dir)
cv2.imwrite(
str(dirs["face_mask"] /f"{video_path.stem}.png"), face_mask)
cv2.imwrite(str(dirs["sep_pose_mask"] /
f"{video_path.stem}.png"), sep_pose_mask)
cv2.imwrite(str(dirs["sep_face_mask"] /
f"{video_path.stem}.png"), sep_face_mask)
cv2.imwrite(str(dirs["sep_lip_mask"] /
f"{video_path.stem}.png"), sep_lip_mask)
else:
images_dir=output_dir/"images"/video_path.stem
audio_path=output_dir/"audios"/f"{video_path.stem}.wav"
_, face_emb, _, _, _=image_processor.preprocess(images_dir)
torch.save(face_emb, str(
dirs["face_emb"] /f"{video_path.stem}.pt"))
audio_emb, _=audio_processor.preprocess(audio_path)
torch.save(audio_emb, str(
dirs["audio_emb"] /f"{video_path.stem}.pt"))
exceptExceptionase:
logging.error(f"Failed to process video {video_path}: {e}")


defprocess_all_videos(input_video_list: List[Path], output_dir: Path, step: int) ->None:
"""
 Process all videos in the input list.

 Args:
 input_video_list (List[Path]): List of video paths to process.
 output_dir (Path): Directory to save the output.
 gpu_status (bool): Whether to use GPU for processing.
 """
face_analysis_model_path="pretrained_models/face_analysis"
landmark_model_path="pretrained_models/face_analysis/models/face_landmarker_v2_with_blendshapes.task"
audio_separator_model_file="pretrained_models/audio_separator/Kim_Vocal_2.onnx"
wav2vec_model_path='pretrained_models/wav2vec/wav2vec2-base-960h'

audio_processor=AudioProcessor(
16000,
25,
wav2vec_model_path,
False,
os.path.dirname(audio_separator_model_file),
os.path.basename(audio_separator_model_file),
os.path.join(output_dir, "vocals"),
 ) ifstep==2elseNone

image_processor=ImageProcessorForDataProcessing(
face_analysis_model_path, landmark_model_path, step)

forvideo_pathintqdm(input_video_list, desc="Processing videos"):
process_single_video(video_path, output_dir,
image_processor, audio_processor, step)


defget_video_paths(source_dir: Path, parallelism: int, rank: int) ->List[Path]:
"""
 Get paths of videos to process, partitioned for parallel processing.

 Args:
 source_dir (Path): Source directory containing videos.
 parallelism (int): Level of parallelism.
 rank (int): Rank for distributed processing.

 Returns:
 List[Path]: List of video paths to process.
 """
video_paths= [itemforiteminsorted(
source_dir.iterdir()) ifitem.is_file() anditem.suffix=='.mp4']
return [video_paths[i] foriinrange(len(video_paths)) ifi%parallelism==rank]


if__name__=="__main__":
parser=argparse.ArgumentParser(
description="Process videos to prepare data for training. Run this script twice with different GPU status parameters."
 )
parser.add_argument("-i", "--input_dir", type=Path,
required=True, help="Directory containing videos")
parser.add_argument("-o", "--output_dir", type=Path,
help="Directory to save results, default is parent dir of input dir")
parser.add_argument("-s", "--step", type=int, default=1,
help="Specify data processing step 1 or 2, you should run 1 and 2 sequently")
parser.add_argument("-p", "--parallelism", default=1,
type=int, help="Level of parallelism")
parser.add_argument("-r", "--rank", default=0, type=int,
help="Rank for distributed processing")

args=parser.parse_args()

ifargs.output_dirisNone:
args.output_dir=args.input_dir.parent

video_path_list=get_video_paths(
args.input_dir, args.parallelism, args.rank)

ifnotvideo_path_list:
logging.warning("No videos to process.")
else:
process_all_videos(video_path_list, args.output_dir, args.step)