Spaces:
Runtime error
Runtime error
| import torch | |
| import cv2 | |
| import sys | |
| import numpy as np | |
| import os | |
| from PIL import Image | |
| # from zdete import Predictor as BboxPredictor | |
| from transformers import Wav2Vec2Model, Wav2Vec2Processor | |
| class MyWav2Vec(): | |
| def __init__(self, model_path, device="cuda"): | |
| super(MyWav2Vec, self).__init__() | |
| self.processor = Wav2Vec2Processor.from_pretrained(model_path) | |
| self.wav2Vec = Wav2Vec2Model.from_pretrained(model_path).to(device) | |
| self.device = device | |
| print("### Wav2Vec model loaded ###") | |
| def forward(self, x): | |
| return self.wav2Vec(x).last_hidden_state | |
| def process(self, x): | |
| return self.processor(x, sampling_rate=16000, return_tensors="pt").input_values.to(self.device) | |
| class AutoFlow(): | |
| def __init__(self, auto_flow_dir, imh=512, imw=512): | |
| super(AutoFlow, self).__init__() | |
| model_dir = auto_flow_dir+'/third_lib/model_zoo/' | |
| cfg_file = model_dir + '/zdete_detector/mobilenet_v1_0.25.yaml' | |
| model_file = model_dir + '/zdete_detector/last_39.pt' | |
| self.bbox_predictor = BboxPredictor(cfg_file, model_file, imgsz=320, conf_thres=0.6, iou_thres=0.2) | |
| self.imh = imh | |
| self.imw = imw | |
| print("### AutoFlow bbox_predictor loaded ###") | |
| def frames_to_face_regions(self, frames, toPIL=True): | |
| # θΎε ₯ζ―bgr numpyζ ΌεΌ | |
| face_region_list = [] | |
| for img in frames: | |
| # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| bbox = self.bbox_predictor.predict(img)[0][0] | |
| xyxy = bbox[:4] | |
| score = bbox[4] | |
| xyxy = np.round(xyxy).astype('int') | |
| rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2] | |
| face_mask = np.zeros((img.shape[0], img.shape[1])).astype('uint8') | |
| face_mask[rb:re,cb:ce] = 255 | |
| face_mask = cv2.resize(face_mask, ((self.imw, self.imh))) | |
| if toPIL: | |
| face_mask = Image.fromarray(face_mask) | |
| face_region_list.append(face_mask) | |
| return face_region_list | |
| def xyxy2x0y0wh(bbox): | |
| x0, y0, x1, y1 = bbox[:4] | |
| return [x0, y0, x1-x0, y1-y0] | |
| def video_to_frame(video_path: str, interval=1, max_frame=None, imh=None, imw=None, is_return_sum=False, is_rgb=False): | |
| vidcap = cv2.VideoCapture(video_path) | |
| success = True | |
| key_frames = [] | |
| sum_frames = None | |
| count = 0 | |
| while success: | |
| success, image = vidcap.read() | |
| if image is not None: | |
| if is_rgb: | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| if imh is not None and imw is not None: | |
| image = img_resize(image, imh=None, imw=None) | |
| if count % interval == 0: | |
| key_frames.append(image) | |
| if is_return_sum: | |
| if sum_frames is None: | |
| sum_frames = image.copy().astype('float32') | |
| else: | |
| sum_frames = sum_frames + image | |
| count += 1 | |
| if max_frame != None: | |
| if count >= max_frame: | |
| break | |
| vidcap.release() | |
| if is_return_sum: | |
| return key_frames, sum_frames | |
| else: | |
| return key_frames | |
| def img_resize(input_img, imh=None, imw=None, max_val=512): | |
| if imh is not None and imw is not None: | |
| width, height = imw, imh | |
| else: | |
| height, width = input_img.shape[0], input_img.shape[1] | |
| if height > width: | |
| ratio = width/height | |
| height = max_val | |
| width = ratio * height | |
| else: | |
| ratio = height/width | |
| width = max_val | |
| height = ratio * width | |
| height = int(round(height/8)*8) | |
| width = int(round(width/8)*8) | |
| input_img = cv2.resize(input_img, (width, height)) | |
| return input_img | |
| def assign_audio_to_frame(audio_input, frame_num): | |
| audio_len = audio_input.shape[0] | |
| audio_per_frame = audio_len / frame_num | |
| audio_to_frame_list = [] | |
| for f_i in range(frame_num): | |
| start_idx = int(round(f_i * audio_per_frame)) | |
| end_idx = int(round((f_i + 1) * audio_per_frame)) | |
| if start_idx >= audio_len: | |
| start_idx = int(round(start_idx - audio_per_frame)) | |
| # print(f"frame_i:{f_i}, start_index:{start_idx}, end_index:{end_idx}, audio_length:{audio_input.shape}") | |
| seg_audio = audio_input[start_idx:end_idx, :] | |
| if type(seg_audio) == np.ndarray: | |
| seg_audio = seg_audio.mean(axis=0, keepdims=True) # B * 20 * 768 | |
| elif torch.is_tensor(seg_audio): | |
| seg_audio = seg_audio.mean(dim=0, keepdim=True) | |
| audio_to_frame_list.append(seg_audio) | |
| if type(seg_audio) == np.ndarray: | |
| audio_to_frames = np.concatenate(audio_to_frame_list, 0) | |
| else: | |
| audio_to_frames = torch.cat(audio_to_frame_list, 0) | |
| return audio_to_frames | |
| def assign_audio_to_frame_new(audio_input, frame_num, pad_frame): | |
| audio_len = audio_input.shape[0] | |
| audio_to_frame_list = [] | |
| for f_i in range(frame_num): | |
| mid_index = int(f_i / frame_num * audio_len) | |
| start_idx = mid_index - pad_frame | |
| end_idx = mid_index + pad_frame + 1 | |
| if start_idx < 0: | |
| start_idx = 0 | |
| end_idx = start_idx + pad_frame * 2 + 1 | |
| if end_idx >= audio_len: | |
| end_idx = audio_len - 1 | |
| start_idx = end_idx - (pad_frame * 2 + 1) | |
| seg_audio = audio_input[None, start_idx:end_idx, :] | |
| audio_to_frame_list.append(seg_audio) | |
| if type(seg_audio) == np.ndarray: | |
| audio_to_frames = np.concatenate(audio_to_frame_list, 0) | |
| else: | |
| audio_to_frames = torch.cat(audio_to_frame_list, 0) | |
| return audio_to_frames | |
| class DotDict(dict): | |
| def __init__(self, *args, **kwargs): | |
| super(DotDict, self).__init__(*args, **kwargs) | |
| def __getattr__(self, key): | |
| value = self[key] | |
| if isinstance(value, dict): | |
| value = DotDict(value) | |
| return value |