| 16 | |
| 17 | |
| 18 | class DataProcessor: |
| 19 | def __init__(self, opt): |
| 20 | self.opt = opt |
| 21 | self.fps = opt.fps |
| 22 | self.sampling_rate = opt.sampling_rate |
| 23 | self.input_size = opt.input_size |
| 24 | |
| 25 | self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False) |
| 26 | |
| 27 | # wav2vec2 audio preprocessor |
| 28 | self.wav2vec_preprocessor = Wav2Vec2FeatureExtractor.from_pretrained(opt.wav2vec_model_path, local_files_only=True) |
| 29 | |
| 30 | # image transform |
| 31 | self.transform = A.Compose([ |
| 32 | A.Resize(height=opt.input_size, width=opt.input_size, interpolation=cv2.INTER_AREA), |
| 33 | A.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5)), |
| 34 | A_pytorch.ToTensorV2(), |
| 35 | ]) |
| 36 | |
| 37 | @torch.no_grad() |
| 38 | def process_img(self, img:np.ndarray) -> np.ndarray: |
| 39 | mult = 360. / img.shape[0] |
| 40 | |
| 41 | resized_img = cv2.resize(img, dsize=(0, 0), fx = mult, fy = mult, interpolation=cv2.INTER_AREA if mult < 1. else cv2.INTER_CUBIC) |
| 42 | bboxes = self.fa.face_detector.detect_from_image(resized_img) |
| 43 | bboxes = [(int(x1 / mult), int(y1 / mult), int(x2 / mult), int(y2 / mult), score) for (x1, y1, x2, y2, score) in bboxes if score > 0.95] |
| 44 | bboxes = bboxes[0] # Just use first bbox |
| 45 | |
| 46 | bsy = int((bboxes[3] - bboxes[1]) / 2) |
| 47 | bsx = int((bboxes[2] - bboxes[0]) / 2) |
| 48 | my = int((bboxes[1] + bboxes[3]) / 2) |
| 49 | mx = int((bboxes[0] + bboxes[2]) / 2) |
| 50 | |
| 51 | bs = int(max(bsy, bsx) * 1.6) |
| 52 | img = cv2.copyMakeBorder(img, bs, bs, bs, bs, cv2.BORDER_CONSTANT, value=0) |
| 53 | my, mx = my + bs, mx + bs # BBox center y, bbox center x |
| 54 | |
| 55 | crop_img = img[my - bs:my + bs,mx - bs:mx + bs] |
| 56 | crop_img = cv2.resize(crop_img, dsize = (self.input_size, self.input_size), interpolation = cv2.INTER_AREA if mult < 1. else cv2.INTER_CUBIC) |
| 57 | return crop_img |
| 58 | |
| 59 | def default_img_loader(self, path) -> np.ndarray: |
| 60 | img = cv2.imread(path) |
| 61 | return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
| 62 | |
| 63 | def default_aud_loader(self, path: str) -> torch.Tensor: |
| 64 | speech_array, sampling_rate = librosa.load(path, sr = self.sampling_rate) |
| 65 | return self.wav2vec_preprocessor(speech_array, sampling_rate = sampling_rate, return_tensors = 'pt').input_values[0] |
| 66 | |
| 67 | |
| 68 | def preprocess(self, ref_path:str, audio_path:str, no_crop:bool) -> dict: |
| 69 | s = self.default_img_loader(ref_path) |
| 70 | if not no_crop: |
| 71 | s = self.process_img(s) |
| 72 | s = self.transform(image=s)['image'].unsqueeze(0) |
| 73 | a = self.default_aud_loader(audio_path).unsqueeze(0) |
| 74 | return {'s': s, 'a': a, 'p': None, 'e': None} |
| 75 | |