Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| import torchvision.transforms as transforms | |
| from transformers import CLIPImageProcessor | |
| import librosa | |
| def process_bbox(bbox, expand_radio, height, width): | |
| """ | |
| raw_vid_path: | |
| bbox: format: x1, y1, x2, y2 | |
| radio: expand radio against bbox size | |
| height,width: source image height and width | |
| """ | |
| def expand(bbox, ratio, height, width): | |
| bbox_h = bbox[3] - bbox[1] | |
| bbox_w = bbox[2] - bbox[0] | |
| expand_x1 = max(bbox[0] - ratio * bbox_w, 0) | |
| expand_y1 = max(bbox[1] - ratio * bbox_h, 0) | |
| expand_x2 = min(bbox[2] + ratio * bbox_w, width) | |
| expand_y2 = min(bbox[3] + ratio * bbox_h, height) | |
| return [expand_x1,expand_y1,expand_x2,expand_y2] | |
| def to_square(bbox_src, bbox_expend, height, width): | |
| h = bbox_expend[3] - bbox_expend[1] | |
| w = bbox_expend[2] - bbox_expend[0] | |
| c_h = (bbox_expend[1] + bbox_expend[3]) / 2 | |
| c_w = (bbox_expend[0] + bbox_expend[2]) / 2 | |
| c = min(h, w) / 2 | |
| c_src_h = (bbox_src[1] + bbox_src[3]) / 2 | |
| c_src_w = (bbox_src[0] + bbox_src[2]) / 2 | |
| s_h, s_w = 0, 0 | |
| if w < h: | |
| d = abs((h - w) / 2) | |
| s_h = min(d, abs(c_src_h-c_h)) | |
| s_h = s_h if c_src_h > c_h else s_h * (-1) | |
| else: | |
| d = abs((h - w) / 2) | |
| s_w = min(d, abs(c_src_w-c_w)) | |
| s_w = s_w if c_src_w > c_w else s_w * (-1) | |
| c_h = (bbox_expend[1] + bbox_expend[3]) / 2 + s_h | |
| c_w = (bbox_expend[0] + bbox_expend[2]) / 2 + s_w | |
| square_x1 = c_w - c | |
| square_y1 = c_h - c | |
| square_x2 = c_w + c | |
| square_y2 = c_h + c | |
| x1, y1, x2, y2 = square_x1, square_y1, square_x2, square_y2 | |
| ww = x2 - x1 | |
| hh = y2 - y1 | |
| cc_x = (x1 + x2)/2 | |
| cc_y = (y1 + y2)/2 | |
| # 1:1 | |
| ww = hh = min(ww, hh) | |
| x1, x2 = round(cc_x - ww/2), round(cc_x + ww/2) | |
| y1, y2 = round(cc_y - hh/2), round(cc_y + hh/2) | |
| return [round(x1), round(y1), round(x2), round(y2)] | |
| bbox_expend = expand(bbox, expand_radio, height=height, width=width) | |
| processed_bbox = to_square(bbox, bbox_expend, height=height, width=width) | |
| return processed_bbox | |
| def get_audio_feature(audio_path, feature_extractor): | |
| audio_input, sampling_rate = librosa.load(audio_path, sr=16000) | |
| assert sampling_rate == 16000 | |
| audio_features = [] | |
| window = 750*640 | |
| for i in range(0, len(audio_input), window): | |
| audio_feature = feature_extractor(audio_input[i:i+window], | |
| sampling_rate=sampling_rate, | |
| return_tensors="pt", | |
| ).input_features | |
| audio_features.append(audio_feature) | |
| audio_features = torch.cat(audio_features, dim=-1) | |
| return audio_features, len(audio_input) // 640 | |
| def image_audio_to_tensor(align_instance, feature_extractor, image_path, audio_path, limit=100, image_size=512, area=1.25): | |
| clip_processor = CLIPImageProcessor() | |
| to_tensor = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) | |
| ]) | |
| mask_to_tensor = transforms.Compose([ | |
| transforms.ToTensor(), | |
| ]) | |
| imSrc_ = Image.open(image_path).convert('RGB') | |
| w, h = imSrc_.size | |
| _, _, bboxes_list = align_instance(np.array(imSrc_)[:,:,[2,1,0]], maxface=True) | |
| if len(bboxes_list) == 0: | |
| return None | |
| bboxSrc = bboxes_list[0] | |
| x1, y1, ww, hh = bboxSrc | |
| x2, y2 = x1 + ww, y1 + hh | |
| mask_img = np.zeros_like(np.array(imSrc_)) | |
| ww, hh = (x2-x1) * area, (y2-y1) * area | |
| center = [(x2+x1)//2, (y2+y1)//2] | |
| x1 = max(center[0] - ww//2, 0) | |
| y1 = max(center[1] - hh//2, 0) | |
| x2 = min(center[0] + ww//2, w) | |
| y2 = min(center[1] + hh//2, h) | |
| mask_img[int(y1):int(y2), int(x1):int(x2)] = 255 | |
| mask_img = Image.fromarray(mask_img) | |
| w, h = imSrc_.size | |
| scale = image_size / min(w, h) | |
| new_w = round(w * scale / 64) * 64 | |
| new_h = round(h * scale / 64) * 64 | |
| if new_h != h or new_w != w: | |
| imSrc = imSrc_.resize((new_w, new_h), Image.LANCZOS) | |
| mask_img = mask_img.resize((new_w, new_h), Image.LANCZOS) | |
| else: | |
| imSrc = imSrc_ | |
| clip_image = clip_processor( | |
| images=imSrc.resize((224, 224), Image.LANCZOS), return_tensors="pt" | |
| ).pixel_values[0] | |
| audio_input, audio_len = get_audio_feature(audio_path, feature_extractor) | |
| audio_len = min(limit, audio_len) | |
| sample = dict( | |
| face_mask=mask_to_tensor(mask_img), | |
| ref_img=to_tensor(imSrc), | |
| clip_images=clip_image, | |
| audio_feature=audio_input[0], | |
| audio_len=audio_len | |
| ) | |
| return sample |