| | 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 |
| | |
| | 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 |
