# encoding: utf-8 import numpy as np import glob import time import cv2 import yaml import os import torch import glob import re import string import copy import json import random import enum import editdistance import pronouncing from torch.utils.data import Dataset import Extractor import options from cvtransforms import * from typing import List, Iterable from helpers import * class CharMap(str, enum.Enum): letters = 'letters' lsr2_text = 'lsr2_text' phonemes = 'phonemes' cmu_phonemes = 'cmu_phonemes' visemes = 'visemes' class Datasets(str, enum.Enum): GRID = 'GRID' LRS2 = 'LRS2' class GridDataset(Dataset): letters = [ ' ', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z' ] lrs2_chars = [ ' ', "'", '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z' ] # phonemes used by the lipnet dataset phonemes = [ ' ', 'AE1', 'AO1', 'D', 'JH', 'Y', 'P', 'AH0', 'OW1', 'G', 'AY1', 'TH', 'IY1', 'CH', 'T', 'AW1', 'F', 'AH1', 'Z', 'R', 'EH1', 'UW1', 'M', 'B', 'W', 'V', 'DH', 'K', 'IH0', 'AA1', 'IH1', 'S', 'EY1', 'N', 'OW0', 'L' ] # full set of phonemes in the CMU phoneme corpus cmu_phonemes = [ ' ', '#', 'AA', 'AA0', 'AA1', 'AA2', 'AE', 'AE0', 'AE1', 'AE2', 'AH', 'AH0', 'AH1', 'AH2', 'AO', 'AO0', 'AO1', 'AO2', 'AW', 'AW0', 'AW1', 'AW2', 'AY', 'AY0', 'AY1', 'AY2', 'B', 'CH', 'D', 'DH', 'EH', 'EH0', 'EH1', 'EH2', 'ER', 'ER0', 'ER1', 'ER2', 'EY', 'EY0', 'EY1', 'EY2', 'F', 'G', 'HH', 'IH', 'IH0', 'IH1', 'IH2', 'IY', 'IY0', 'IY1', 'IY2', 'JH', 'K', 'L', 'M', 'N', 'NG', 'OW', 'OW0', 'OW1', 'OW2', 'OY', 'OY0', 'OY1', 'OY2', 'P', 'R', 'S', 'SH', 'T', 'TH', 'UH', 'UH0', 'UH1', 'UH2', 'UW', 'UW0', 'UW1', 'UW2', 'V', 'W', 'Y', 'Z', 'ZH' ] phoneme_chars = map_phonemes(phonemes) cmu_phoneme_chars = map_phonemes(cmu_phonemes) def __init__( self, video_path, alignments_dir, phonemes_dir, file_list, vid_pad, image_dir, txt_pad, phase, shared_dict=None, char_map=CharMap.letters, base_dir='', frame_doubling=False, sample_all_props=False ): self.base_dir = base_dir self.sample_all_props = sample_all_props self.image_dir = os.path.join(base_dir, image_dir) self.alignments_dir = os.path.join(base_dir, alignments_dir) self.phonemes_dir = os.path.join(base_dir, phonemes_dir) self.frame_doubling = frame_doubling if type(file_list) is str: file_list = os.path.join(base_dir, file_list) # print('FILE_LIST', file_list, base_dir) file_list = open(file_list, 'r').readlines() self.shared_dict = shared_dict self.char_map = char_map self.vid_pad = vid_pad self.txt_pad = txt_pad self.phase = phase self.videos = [ os.path.join(video_path, line.strip()) for line in file_list ] self.data = [] for vid in self.videos: items = vid.split(os.path.sep) if len(items) < 2: print('BAD VID ITEM', items) raise ValueError speaker_name, filename = items[-2], items[-1] self.data.append((vid, speaker_name, filename)) def _fetch_anno_path(self, spk, basename): return self.fetch_anno_path( spk=spk, basename=basename, char_map=self.char_map ) @classmethod def text_to_phonemes( cls, text, as_str=True, char_map=CharMap.phonemes ): sentence_phonemes = [] words = text.upper().strip().split(' ') for word in words: word_phonemes = pronouncing.phones_for_word(word)[0] word_phonemes = word_phonemes.split(' ') sentence_phonemes.extend(word_phonemes) sentence_phonemes.append(' ') if sentence_phonemes[-1] == ' ': sentence_phonemes = sentence_phonemes[:-1] if as_str: return cls.stringify(sentence_phonemes, char_map=char_map) else: return sentence_phonemes def fetch_anno_path(self, spk, basename, char_map): if char_map == CharMap.letters: align_path_name = os.path.join( self.alignments_dir, spk, basename + '.align' ) return align_path_name elif char_map == CharMap.lsr2_text: align_path_name = os.path.join( self.alignments_dir, spk, basename + '.txt' ) return align_path_name elif char_map == CharMap.phonemes: phonemes_path_name = os.path.join( self.phonemes_dir, spk, basename + '.align' ) return phonemes_path_name elif char_map == CharMap.cmu_phonemes: phonemes_path_name = os.path.join( self.phonemes_dir, spk, basename + '.txt' ) return phonemes_path_name else: raise NotImplementedError def fetch_anno_text(self, spk, basename, char_map: CharMap): return self.load_anno_text(self.fetch_anno_path( spk, basename, char_map=char_map ), char_map=char_map) def __getitem__(self, idx): (vid, spk, name) = self.data[idx] return self.load_sample( video_name=vid, speaker_name=spk, filename=name ) def load_random_sample(self, char_map=None): (vid, spk, name) = random.choice(self.data) return self.load_sample( video_name=vid, speaker_name=spk, filename=name, char_map=char_map ) def load_sample( self, video_name, speaker_name, filename, char_map=None ): if char_map is None: char_map = self.char_map if self.sample_all_props: char_map = all vid = self.load_vid(video_name) if self.frame_doubling: vid = np.repeat(vid, repeats=2, axis=0) basename, _ = os.path.splitext(filename) # print('SPK_NAME', (spk, name, anno_path)) txt_results, phoneme_results = {}, {} cmu_phoneme_results = {} if (char_map is all) or (char_map == CharMap.letters): txt_anno, txt_anno_arr = self.fetch_anno_text( speaker_name, basename, char_map=CharMap.letters ) txt_anno_arr_len = txt_anno_arr.shape[0] txt_anno_arr = self._padding(txt_anno_arr, self.txt_pad) assert not np.isnan(txt_anno_arr).any() txt_anno += [' '] * (options.txt_padding - len(txt_anno)) txt_results = kwargify( txt=torch.LongTensor(txt_anno_arr), txt_len=txt_anno_arr_len, txt_anno=txt_anno ) if (char_map is all) or (char_map == CharMap.phonemes): phoneme_anno, phoneme_anno_arr = self.fetch_anno_text( speaker_name, basename, char_map=CharMap.phonemes ) phoneme_anno_arr_len = phoneme_anno_arr.shape[0] phoneme_anno_arr = self._padding( phoneme_anno_arr, self.txt_pad ) assert not np.isnan(phoneme_anno_arr_len).any() phoneme_results = kwargify( phonemes=torch.LongTensor(phoneme_anno_arr), phonemes_len=phoneme_anno_arr_len, ) elif (char_map is all) or (char_map == CharMap.cmu_phonemes): phoneme_anno, phoneme_anno_arr = self.fetch_anno_text( speaker_name, basename, char_map=CharMap.cmu_phonemes ) phoneme_anno_arr_len = phoneme_anno_arr.shape[0] phoneme_anno_arr = self._padding( phoneme_anno_arr, self.txt_pad ) assert not np.isnan(phoneme_anno_arr_len).any() cmu_phoneme_results = kwargify( cmu_phonemes=torch.LongTensor(phoneme_anno_arr), cmu_phonemes_len=phoneme_anno_arr_len, ) if self.phase == 'train': vid = HorizontalFlip(vid) vid = ColorNormalize(vid) vid_len = vid.shape[0] vid = self._padding(vid, self.vid_pad) """ if vid_len <= anno_len * 2: raise ValueError(f'CTC INVALID: {self.data[idx]}') """ assert not np.isnan(vid).any() return kwargify( vid=torch.FloatTensor(vid.transpose(3, 0, 1, 2)), vid_len=vid_len, **txt_results, **phoneme_results, **cmu_phoneme_results ) def __len__(self): return len(self.data) @staticmethod def serialize(data: np.ndarray): return torch.from_numpy(data.astype(np.uint8)) @staticmethod def deserialize(data: torch.Tensor): return data.numpy().astype(np.float16) @staticmethod def process_vid(video_path: str, to_tensor=True): frames = Extractor.extract_frames( video_path, recycle_landmarks=True, use_gpu=True ) frames = [f for f in frames if f is not None] array = list(filter(lambda im: im is not None, frames)) array = [ cv2.resize(im, (128, 64), interpolation=cv2.INTER_LANCZOS4) for im in array ] array = np.stack(array, axis=0).astype(np.float16) vid = ColorNormalize(array) if to_tensor: vid = torch.FloatTensor(vid.transpose(3, 0, 1, 2)) return vid def load_vid(self, video_path: str) -> np.ndarray: return self._load_vid(video_path, cache=False) def _load_vid(self, video_path: str, cache=True) -> np.ndarray: if cache and self.shared_dict is not None: if video_path in self.shared_dict: return self.deserialize( self.shared_dict[video_path] ) # print('LOAD_DIR', video_path) base_filename = os.path.basename(video_path) basename, _ = os.path.splitext(base_filename) speaker_dir = os.path.basename(os.path.dirname(video_path)) image_dir = f'{self.image_dir}/{speaker_dir}/{basename}' files = os.listdir(image_dir) files = list(filter(lambda file: file.find('.jpg') != -1, files)) files = sorted(files, key=lambda file: int(os.path.splitext(file)[0])) array = [cv2.imread(os.path.join(image_dir, file)) for file in files] array = list(filter(lambda im: im is not None, array)) array = [ cv2.resize(im, (128, 64), interpolation=cv2.INTER_LANCZOS4) for im in array ] try: array = np.stack(array, axis=0).astype(np.float16) except ValueError as e: print(f'BAD VIDEO PATH: {video_path}') raise e if cache and self.shared_dict is not None: # print('SD >>') serialized_data = self.serialize(array) serialized_data.share_memory_() self.shared_dict[video_path] = serialized_data # print('SD <<') return array @classmethod def load_anno(cls, name, char_map): return cls.load_anno_text(name, char_map)[1] @classmethod def load_anno_text(cls, name, char_map): # print('ANNOTATION_NAME', name) txt = cls.load_sentence(name, char_map=char_map) indices = cls.txt2arr(txt, 1, char_map=char_map) # print('TXT', txt) return txt, indices def _load_anno(self, name): return self.load_anno(name, self.char_map) @classmethod def load_sentence(cls, name, char_map=CharMap.letters) -> List[str]: with open(name, 'r') as f: if char_map == CharMap.letters: lines = [line.strip().split(' ') for line in f.readlines()] txt = [line[2] for line in lines] txt = list(filter( lambda s: not s.upper() in ['SIL', 'SP'], txt )) all_chars = list(' '.join(txt)) all_chars = [char.upper() for char in all_chars] return all_chars elif char_map == CharMap.lsr2_text: text_line = f.readlines()[0] text_line = text_line[5:].strip() all_chars = [char.upper() for char in text_line] return all_chars elif char_map in (CharMap.phonemes, CharMap.cmu_phonemes): all_chars = [] for line in f.readlines(): word_phonemes = line.strip().split(' ') all_chars.extend(word_phonemes) all_chars.append(' ') if all_chars[-1] == ' ': all_chars = all_chars[:-1] return all_chars else: raise ValueError(f'BAD CHAR MAP {char_map}') @classmethod def load_str_sentence(cls, name, char_map=CharMap.letters) -> str: chars_seq = cls.load_sentence(name=name, char_map=char_map) return cls.stringify(chars_seq, char_map=char_map) @staticmethod def tokenize_text(text: str, word_tokenize=False) -> List[str]: """ :param text: :param word_tokenize: whether to tokenize into words or individual characters :return: """ if word_tokenize: return text.split(' ') else: return list(text) @staticmethod def tokenize_phonemes(text: str, word_tokenize=False) -> List[str]: """ :param text: :param word_tokenize: whether to tokenize into words or individual phonemes example: text = 'S-EH1-T G-R-IY1-N IH0-N EH1-L S-IH1-K-S AH0-G-EH1-N' word-level tokens: ['S-EH1-T', 'G-R-IY1-N', 'IH0-N', 'EH1-L', 'S-IH1-K-S', 'AH0-G-EH1-N'] phoneme-level tokens: ['S', 'EH1', 'T', ' ', 'G', 'R', 'IY1', 'N', ' ', 'IH0', 'N', ' ', 'EH1', 'L', ' ', 'S', 'IH1', 'K', 'S', ' ', 'AH0', 'G', 'EH1', 'N'] :return: """ if word_tokenize: return text.split(' ') else: words = text.split(' ') phonemes = [] for word in words: assert not word.startswith('-') assert not word.endswith('-') phonemes.extend(word.split('-')) phonemes.append(' ') if phonemes[-1] == ' ': phonemes = phonemes[:-1] return phonemes @staticmethod def _padding(array, length): array = [array[_] for _ in range(array.shape[0])] size = array[0].shape for i in range(length - len(array)): array.append(np.zeros(size)) return np.stack(array, axis=0) @classmethod def txt2arr(cls, txt, start, char_map=CharMap.letters): arr = [] if char_map == CharMap.letters: for char in list(txt): arr.append(cls.letters.index(char) + start) elif char_map == CharMap.phonemes: # print('TXT', txt) for phoneme in txt: arr.append(cls.phonemes.index(phoneme) + start) elif char_map == CharMap.cmu_phonemes: # print('TXT', txt) for phoneme in txt: arr.append(cls.cmu_phonemes.index(phoneme) + start) elif char_map == CharMap.visemes: raise NotImplementedError else: raise ValueError(f'BAD CHAR MAP: {char_map}') return np.array(arr) def arr2txt(self, arr, start, char_map=None): char_map = self.char_map if char_map is None else char_map return self._arr2txt(arr, start, char_map=char_map) @classmethod def _arr2txt(cls, arr, start, char_map=CharMap.letters): txt = [] for n in arr: if n >= start: if char_map == CharMap.letters: txt.append(cls.letters[n - start]) elif char_map == CharMap.phonemes: txt.append(cls.phonemes[n - start]) elif char_map == CharMap.cmu_phonemes: txt.append(cls.cmu_phonemes[n - start]) elif char_map == CharMap.visemes: raise NotImplementedError else: raise ValueError(f'BAD CHAR MAP: {char_map}') return cls.stringify(txt, char_map) def get_char_mapping(self): return self.char_mapping(self.char_map) @classmethod def char_mapping(cls, char_map): if char_map == CharMap.letters: return cls.letters elif char_map == CharMap.phonemes: return cls.phonemes elif char_map == CharMap.cmu_phonemes: return cls.cmu_phonemes elif char_map == CharMap.visemes: raise NotImplementedError else: raise ValueError(f'BAD CHAR MAP: {char_map}') def ctc_decode(self, y): y = y.argmax(-1) return [ self.ctc_arr2txt(y[_], start=1) for _ in range(y.size(0)) ] def ctc_decode_indices(self, y): y = y.argmax(-1) return [ self.ctc_arr2txt_indices(y[_], start=1)[1] for _ in range(y.size(0)) ] def ctc_arr2txt(self, *args, **kwargs): sentence, indices = self.ctc_arr2txt_pair(*args, **kwargs) return sentence def ctc_arr2txt_pair( self, arr, start, char_map=None, filter_previous=True ): """ converts token indices into a string sentence :param arr: array of token indices :param start: number of special characters in character set :param char_map: character set to use for tokenization :param filter_previous: if True, removes consecutive occurrences of an index / token e.g. THREE becomes THRE, SOON becomes SON :return: """ sentence, indices = self.ctc_arr2txt_indices( arr=arr, start=start, char_map=char_map, filter_previous=filter_previous ) return sentence, indices def ctc_arr2txt_indices( self, arr, start, char_map=None, filter_previous=True ): """ converts token indices into a string sentence and indices of tokens taken along arr :param arr: array of token indices :param start: number of special characters in character set :param char_map: character set to use for tokenization :param filter_previous: if True, removes consecutive occurrences of an index / token e.g. THREE becomes THRE, SOON becomes SON :return: """ if char_map is None: char_map = self.char_map previous = -1 txt, indices = [], [] char_mapping = self.char_mapping(char_map) for k, n in enumerate(arr): check_consecutive = ( not filter_previous or previous != n ) if n >= start: has_empty_char = ( len(txt) > 0 and txt[-1] == ' ' and char_mapping[n - start] == ' ' ) if not has_empty_char and check_consecutive: txt.append(char_mapping[n - start]) indices.append(k) previous = n sentence = self.stringify(txt, char_map) return sentence, indices @staticmethod def stringify(txt, char_map): if char_map in (CharMap.letters, CharMap.lsr2_text): return ''.join(txt).strip() elif char_map in (CharMap.phonemes, CharMap.cmu_phonemes): sentence = '-'.join(txt).strip() sentence = sentence.replace('- ', ' ') sentence = sentence.replace(' -', ' ') if sentence.endswith('-'): sentence = sentence[:-1] if sentence.startswith('-'): sentence = sentence[1:] return sentence else: raise NotImplementedError def _map_chars(self, chars: str): return self.map_chars(chars, char_map=self.char_map) @classmethod def map_chars(cls, chars: str, char_map: CharMap): # map a string containing multi-character # phonemes like AE1 to a single character if char_map == CharMap.letters: return chars elif char_map in (CharMap.phonemes, CharMap.cmu_phonemes): if char_map == CharMap.phonemes: phonemes_arr = cls.phonemes char_phonemes_arr = cls.phonemes elif char_map == CharMap.cmu_phonemes: phonemes_arr = cls.cmu_phonemes char_phonemes_arr = cls.cmu_phoneme_chars else: raise ValueError(f'BAD CHAR MAP {char_map}') words = chars.split(' ') char_phonemes = '' for word in words: phonemes = word.split('-') phonemes = [ phoneme for phoneme in phonemes if phoneme.strip() != '' ] for phoneme in phonemes: char_phonemes += char_phonemes_arr[ phonemes_arr.index(phoneme) ] char_phonemes += ' ' return char_phonemes elif char_map == CharMap.visemes: raise NotImplementedError else: raise ValueError(f'BAD CHAR MAP: {char_map}') @classmethod def map_char_lists( cls, char_lists: Iterable[str], char_map: CharMap ): return [cls.map_chars( char_seq, char_map=char_map ) for char_seq in char_lists] def wer(self, raw_predict, raw_truth): return self.get_wer( raw_predict, raw_truth, char_map=self.char_map ) @classmethod def get_wer(cls, raw_predict, raw_truth, char_map: CharMap): assert isinstance(raw_predict, Iterable) assert isinstance(raw_truth, Iterable) predict = cls.map_char_lists(raw_predict, char_map=char_map) truth = cls.map_char_lists(raw_truth, char_map=char_map) # print('WER', raw_truth, raw_predict) word_pairs = [ (p[0].split(' '), p[1].split(' ')) for p in zip(predict, truth) ] wer = [ 1.0 * editdistance.eval(p[0], p[1])/len(p[1]) for p in word_pairs ] return wer def cer(self, raw_predict, raw_truth): return self.get_cer( raw_predict, raw_truth, char_map=self.char_map ) @classmethod def get_cer(cls, raw_predict, raw_truth, char_map: CharMap): assert isinstance(raw_predict, Iterable) assert isinstance(raw_truth, Iterable) predict = cls.map_char_lists(raw_predict, char_map=char_map) truth = cls.map_char_lists(raw_truth, char_map=char_map) cer = [ 1.0 * editdistance.eval(p[0], p[1]) / len(p[1]) for p in zip(predict, truth) ] return cer