Delete src

src/dataset/test_preprocess.py

@@ -1,155 +0,0 @@
-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