File size: 2,630 Bytes
3a3c68a e1cca9e 3a3c68a e1cca9e bd9d6b3 e1cca9e 3a3c68a |
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 |
import os
from typing import List
import yaml
class GridLoader(object):
def __init__(self):
with open('config.yml', 'r') as config_file_obj:
yaml_config = yaml.safe_load(config_file_obj)
self.dataset_config = yaml_config['datasets']
self.video_dir = self.dataset_config['video_dir']
self.alignment_dir = self.dataset_config['alignments_dir']
self.usable_video_filepaths = None
def load_videos(
self, verbose=False, blacklist=frozenset({})
) -> List[str]:
"""
:param verbose:
whether to show logs
(currently displays numbers of videos with alignment loaded)
:param blacklist:
set of filepaths to exclude from training
:return:
"""
usable_video_filepaths = []
videos_without_alignment = []
for speaker_no in range(1, 35):
speaker_dirname = f's{speaker_no}'
speaker_dir = os.path.join(self.video_dir, speaker_dirname)
if not os.path.exists(speaker_dir):
# speaker does not exist (its just s21 right now)
continue
video_filenames = os.listdir(speaker_dir)
for video_filename in video_filenames:
if not video_filename.endswith('.mpg'):
continue
# get name of file without the extension
base_name = os.path.splitext(video_filename)[0]
video_path = os.path.join(
self.video_dir, speaker_dirname, f'{base_name}.mpg'
)
if video_path in blacklist:
continue
alignment_path = os.path.join(
self.alignment_dir, speaker_dirname, f'{base_name}.align'
)
if os.path.exists(alignment_path):
# don't include video if it has no corresponding
# alignment path
usable_video_filepaths.append(video_path)
else:
videos_without_alignment.append(alignment_path)
if verbose:
num_usable_videos = len(usable_video_filepaths)
num_unusable_videos = len(videos_without_alignment)
# print(videos_without_alignment)
print(f'videos with alignment: {num_usable_videos}')
print(f'videos without alignment: {num_unusable_videos}')
self.usable_video_filepaths = usable_video_filepaths
return usable_video_filepaths
if __name__ == '__main__':
loader = GridLoader()
loader.load_videos(True) |