import warnings warnings.filterwarnings("ignore") import os import argparse import face_alignment import torch import torchaudio import numpy as np import cv2 from PIL import Image, ImageDraw from moviepy import * from collections import deque from skimage import transform as tf import yaml from look2hear.models import Dolphin from look2hear.datas.transform import get_preprocessing_pipelines from face_detection_utils import detect_faces # -- Landmark interpolation: def linear_interpolate(landmarks, start_idx, stop_idx): start_landmarks = landmarks[start_idx] stop_landmarks = landmarks[stop_idx] delta = stop_landmarks - start_landmarks for idx in range(1, stop_idx-start_idx): landmarks[start_idx+idx] = start_landmarks + idx/float(stop_idx-start_idx) * delta return landmarks # -- Face Transformation def warp_img(src, dst, img, std_size): tform = tf.estimate_transform('similarity', src, dst) # find the transformation matrix warped = tf.warp(img, inverse_map=tform.inverse, output_shape=std_size) # wrap the frame image warped = warped * 255 # note output from wrap is double image (value range [0,1]) warped = warped.astype('uint8') return warped, tform def apply_transform(transform, img, std_size): warped = tf.warp(img, inverse_map=transform.inverse, output_shape=std_size) warped = warped * 255 # note output from wrap is double image (value range [0,1]) warped = warped.astype('uint8') return warped # -- Crop def cut_patch(img, landmarks, height, width, threshold=5): center_x, center_y = np.mean(landmarks, axis=0) if center_y - height < 0: center_y = height if center_y - height < 0 - threshold: raise Exception('too much bias in height') if center_x - width < 0: center_x = width if center_x - width < 0 - threshold: raise Exception('too much bias in width') if center_y + height > img.shape[0]: center_y = img.shape[0] - height if center_y + height > img.shape[0] + threshold: raise Exception('too much bias in height') if center_x + width > img.shape[1]: center_x = img.shape[1] - width if center_x + width > img.shape[1] + threshold: raise Exception('too much bias in width') cutted_img = np.copy(img[ int(round(center_y) - round(height)): int(round(center_y) + round(height)), int(round(center_x) - round(width)): int(round(center_x) + round(width))]) return cutted_img # -- RGB to GRAY def convert_bgr2gray(data): return np.stack([cv2.cvtColor(_, cv2.COLOR_BGR2GRAY) for _ in data], axis=0) def save2npz(filename, data=None): assert data is not None, "data is {}".format(data) if not os.path.exists(os.path.dirname(filename)): os.makedirs(os.path.dirname(filename)) np.savez_compressed(filename, data=data) def read_video(filename): """Read video frames using MoviePy for better compatibility""" try: video_clip = VideoFileClip(filename) for frame in video_clip.iter_frames(): # Convert RGB to BGR to match cv2 format frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) yield frame_bgr video_clip.close() except Exception as e: print(f"Error reading video {filename}: {e}") return def face2head(boxes, scale=1.5): new_boxes = [] for box in boxes: width = box[2] - box[0] height= box[3] - box[1] width_center = (box[2] + box[0]) / 2 height_center = (box[3] + box[1]) / 2 square_width = int(max(width, height) * scale) new_box = [width_center - square_width/2, height_center - square_width/2, width_center + square_width/2, height_center + square_width/2] new_boxes.append(new_box) return new_boxes def bb_intersection_over_union(boxA, boxB): # determine the (x, y)-coordinates of the intersection rectangle xA = max(boxA[0], boxB[0]) yA = max(boxA[1], boxB[1]) xB = min(boxA[2], boxB[2]) yB = min(boxA[3], boxB[3]) # compute the area of intersection rectangle interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1) # compute the area of both the prediction and ground-truth # rectangles boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1) boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1) # compute the intersection over union by taking the intersection # area and dividing it by the sum of prediction + ground-truth # areas - the interesection area iou = interArea / float(boxAArea + boxBArea - interArea) # return the intersection over union value return iou def detectface(video_input_path, output_path, detect_every_N_frame, scalar_face_detection, number_of_speakers): device = torch.device('cuda' if torch.cuda.get_device_name() else 'cpu') print('Running on device: {}'.format(device)) os.makedirs(os.path.join(output_path, 'faces'), exist_ok=True) os.makedirs(os.path.join(output_path, 'landmark'), exist_ok=True) landmarks_dic = {} faces_dic = {} boxes_dic = {} for i in range(number_of_speakers): landmarks_dic[i] = [] faces_dic[i] = [] boxes_dic[i] = [] video_clip = VideoFileClip(video_input_path) print("Video statistics: ", video_clip.w, video_clip.h, (video_clip.w, video_clip.h), video_clip.fps) frames = [Image.fromarray(frame) for frame in video_clip.iter_frames()] print('Number of frames in video: ', len(frames)) video_clip.close() fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False) for i, frame in enumerate(frames): print('\rTracking frame: {}'.format(i + 1), end='') # Detect faces every N frames if i % detect_every_N_frame == 0: frame_array = np.array(frame) detected_boxes, _ = detect_faces( frame_array, threshold=0.9, allow_upscaling=False, ) if detected_boxes is None or len(detected_boxes) == 0: detected_boxes, _ = detect_faces( frame_array, threshold=0.7, allow_upscaling=True, ) if detected_boxes is not None and len(detected_boxes) > 0: detected_boxes = detected_boxes[:number_of_speakers] detected_boxes = face2head(detected_boxes, scalar_face_detection) else: detected_boxes = [] # Process the detection results if i == 0: # First frame - initialize tracking if len(detected_boxes) < number_of_speakers: raise ValueError(f"First frame must detect at least {number_of_speakers} faces, but only found {len(detected_boxes)}") # Assign first detections to speakers in order for j in range(number_of_speakers): box = detected_boxes[j] face = frame.crop((box[0], box[1], box[2], box[3])).resize((224,224)) preds = fa.get_landmarks(np.array(face)) if preds is None: raise ValueError(f"Face landmarks not detected in initial frame for speaker {j}") faces_dic[j].append(face) landmarks_dic[j].append(preds) boxes_dic[j].append(box) else: # For subsequent frames, match detected boxes to speakers matched_speakers = set() speaker_boxes = [None] * number_of_speakers # Match each detected box to the most likely speaker for box in detected_boxes: iou_scores = [] for speaker_id in range(number_of_speakers): if speaker_id in matched_speakers: iou_scores.append(-1) # Already matched else: last_box = boxes_dic[speaker_id][-1] iou_score = bb_intersection_over_union(box, last_box) iou_scores.append(iou_score) if max(iou_scores) > 0: # Valid match found best_speaker = iou_scores.index(max(iou_scores)) speaker_boxes[best_speaker] = box matched_speakers.add(best_speaker) # Process each speaker for speaker_id in range(number_of_speakers): if speaker_boxes[speaker_id] is not None: # Use detected box box = speaker_boxes[speaker_id] else: # Use previous box for this speaker box = boxes_dic[speaker_id][-1] # Extract face and landmarks face = frame.crop((box[0], box[1], box[2], box[3])).resize((224,224)) preds = fa.get_landmarks(np.array(face)) if preds is None: # Use previous landmarks if detection fails preds = landmarks_dic[speaker_id][-1] faces_dic[speaker_id].append(face) landmarks_dic[speaker_id].append(preds) boxes_dic[speaker_id].append(box) # Verify all speakers have same number of frames frame_counts = [len(boxes_dic[s]) for s in range(number_of_speakers)] print(f"\nFrame counts per speaker: {frame_counts}") assert all(count == len(frames) for count in frame_counts), f"Inconsistent frame counts: {frame_counts}" # Continue with saving videos and landmarks... for s in range(number_of_speakers): frames_tracked = [] for i, frame in enumerate(frames): frame_draw = frame.copy() draw = ImageDraw.Draw(frame_draw) draw.rectangle(boxes_dic[s][i], outline=(255, 0, 0), width=6) frames_tracked.append(frame_draw) # Save tracked video tracked_frames = [np.array(frame) for frame in frames_tracked] if tracked_frames: tracked_clip = ImageSequenceClip(tracked_frames, fps=25.0) tracked_video_path = os.path.join(output_path, 'video_tracked' + str(s+1) + '.mp4') tracked_clip.write_videofile(tracked_video_path, codec='libx264', audio=False, logger=None) tracked_clip.close() # Save landmarks for i in range(number_of_speakers): save2npz(os.path.join(output_path, 'landmark', 'speaker' + str(i+1)+'.npz'), data=landmarks_dic[i]) # Save face video face_frames = [np.array(frame) for frame in faces_dic[i]] if face_frames: face_clip = ImageSequenceClip(face_frames, fps=25.0) face_video_path = os.path.join(output_path, 'faces', 'speaker' + str(i+1) + '.mp4') face_clip.write_videofile(face_video_path, codec='libx264', audio=False, logger=None) face_clip.close() # Output video path parts = video_input_path.split('/') video_name = parts[-1][:-4] if not os.path.exists(os.path.join(output_path, 'filename_input')): os.mkdir(os.path.join(output_path, 'filename_input')) csvfile = open(os.path.join(output_path, 'filename_input', str(video_name) + '.csv'), 'w') for i in range(number_of_speakers): csvfile.write('speaker' + str(i+1)+ ',0\n') csvfile.close() return os.path.join(output_path, 'filename_input', str(video_name) + '.csv') def crop_patch(mean_face_landmarks, video_pathname, landmarks, window_margin, start_idx, stop_idx, crop_height, crop_width, STD_SIZE=(256, 256)): """Crop mouth patch :param str video_pathname: pathname for the video_dieo :param list landmarks: interpolated landmarks """ stablePntsIDs = [33, 36, 39, 42, 45] frame_idx = 0 frame_gen = read_video(video_pathname) while True: try: frame = frame_gen.__next__() ## -- BGR except StopIteration: break if frame_idx == 0: q_frame, q_landmarks = deque(), deque() sequence = [] q_landmarks.append(landmarks[frame_idx]) q_frame.append(frame) if len(q_frame) == window_margin: smoothed_landmarks = np.mean(q_landmarks, axis=0) cur_landmarks = q_landmarks.popleft() cur_frame = q_frame.popleft() # -- affine transformation trans_frame, trans = warp_img( smoothed_landmarks[stablePntsIDs, :], mean_face_landmarks[stablePntsIDs, :], cur_frame, STD_SIZE) trans_landmarks = trans(cur_landmarks) # -- crop mouth patch sequence.append( cut_patch( trans_frame, trans_landmarks[start_idx:stop_idx], crop_height//2, crop_width//2,)) if frame_idx == len(landmarks)-1: #deal with corner case with video too short if len(landmarks) < window_margin: smoothed_landmarks = np.mean(q_landmarks, axis=0) cur_landmarks = q_landmarks.popleft() cur_frame = q_frame.popleft() # -- affine transformation trans_frame, trans = warp_img(smoothed_landmarks[stablePntsIDs, :], mean_face_landmarks[stablePntsIDs, :], cur_frame, STD_SIZE) trans_landmarks = trans(cur_landmarks) # -- crop mouth patch sequence.append(cut_patch( trans_frame, trans_landmarks[start_idx:stop_idx], crop_height//2, crop_width//2,)) while q_frame: cur_frame = q_frame.popleft() # -- transform frame trans_frame = apply_transform( trans, cur_frame, STD_SIZE) # -- transform landmarks trans_landmarks = trans(q_landmarks.popleft()) # -- crop mouth patch sequence.append( cut_patch( trans_frame, trans_landmarks[start_idx:stop_idx], crop_height//2, crop_width//2,)) return np.array(sequence) frame_idx += 1 return None def landmarks_interpolate(landmarks): """Interpolate landmarks param list landmarks: landmarks detected in raw videos """ valid_frames_idx = [idx for idx, _ in enumerate(landmarks) if _ is not None] if not valid_frames_idx: return None for idx in range(1, len(valid_frames_idx)): if valid_frames_idx[idx] - valid_frames_idx[idx-1] == 1: continue else: landmarks = linear_interpolate(landmarks, valid_frames_idx[idx-1], valid_frames_idx[idx]) valid_frames_idx = [idx for idx, _ in enumerate(landmarks) if _ is not None] # -- Corner case: keep frames at the beginning or at the end failed to be detected. if valid_frames_idx: landmarks[:valid_frames_idx[0]] = [landmarks[valid_frames_idx[0]]] * valid_frames_idx[0] landmarks[valid_frames_idx[-1]:] = [landmarks[valid_frames_idx[-1]]] * (len(landmarks) - valid_frames_idx[-1]) valid_frames_idx = [idx for idx, _ in enumerate(landmarks) if _ is not None] assert len(valid_frames_idx) == len(landmarks), "not every frame has landmark" return landmarks def crop_mouth(video_direc, landmark_direc, filename_path, save_direc, convert_gray=False, testset_only=False): lines = open(filename_path).read().splitlines() lines = list(filter(lambda x: 'test' in x, lines)) if testset_only else lines for filename_idx, line in enumerate(lines): filename, person_id = line.split(',') print('idx: {} \tProcessing.\t{}'.format(filename_idx, filename)) video_pathname = os.path.join(video_direc, filename+'.mp4') landmarks_pathname = os.path.join(landmark_direc, filename+'.npz') dst_pathname = os.path.join( save_direc, filename+'.npz') # if os.path.exists(dst_pathname): # continue multi_sub_landmarks = np.load(landmarks_pathname, allow_pickle=True)['data'] landmarks = [None] * len(multi_sub_landmarks) for frame_idx in range(len(landmarks)): try: #landmarks[frame_idx] = multi_sub_landmarks[frame_idx][int(person_id)]['facial_landmarks'] #original for LRW landmarks[frame_idx] = multi_sub_landmarks[frame_idx][int(person_id)] #VOXCELEB2 except (IndexError, TypeError): continue # -- pre-process landmarks: interpolate frames not being detected. preprocessed_landmarks = landmarks_interpolate(landmarks) if not preprocessed_landmarks: continue # -- crop mean_face_landmarks = np.load('assets/20words_mean_face.npy') sequence = crop_patch(mean_face_landmarks, video_pathname, preprocessed_landmarks, 12, 48, 68, 96, 96) assert sequence is not None, "cannot crop from {}.".format(filename) # -- save data = convert_bgr2gray(sequence) if convert_gray else sequence[...,::-1] save2npz(dst_pathname, data=data) def convert_video_fps(input_file, output_file, target_fps=25): """Convert video to target FPS using moviepy""" video = VideoFileClip(input_file) video_fps = video.fps if video_fps != target_fps: video.write_videofile( output_file, fps=target_fps, codec='libx264', audio_codec='aac', temp_audiofile='temp-audio.m4a', remove_temp=True, ) else: # If already at target fps, just copy import shutil shutil.copy2(input_file, output_file) video.close() print(f'Video has been converted to {target_fps} fps and saved to {output_file}') def extract_audio(video_file, audio_output_file, sample_rate=16000): """Extract audio from video using moviepy""" video = VideoFileClip(video_file) audio = video.audio # Save audio with specified sample rate audio.write_audiofile(audio_output_file, fps=sample_rate, nbytes=2, codec='pcm_s16le') video.close() audio.close() def merge_video_audio(video_file, audio_file, output_file): """Merge video and audio using moviepy""" video = VideoFileClip(video_file) audio = AudioFileClip(audio_file) # Attach audio (MoviePy v2 renamed set_audio to with_audio) set_audio_fn = getattr(video, "set_audio", None) if callable(set_audio_fn): final_video = set_audio_fn(audio) else: with_audio_fn = getattr(video, "with_audio", None) if not callable(with_audio_fn): video.close() audio.close() raise AttributeError("VideoFileClip object lacks both set_audio and with_audio methods") final_video = with_audio_fn(audio) # Write the result final_video.write_videofile(output_file, codec='libx264', audio_codec='aac', temp_audiofile='temp-audio.m4a', remove_temp=True) # Clean up video.close() audio.close() final_video.close() def process_video(input_file, output_path, number_of_speakers=2, detect_every_N_frame=8, scalar_face_detection=1.5, config_path="checkpoints/vox2/conf.yml", cuda_device=None): """Main processing function for video speaker separation""" # Set CUDA device if specified if cuda_device is not None: os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_device) # Create output directory os.makedirs(output_path, exist_ok=True) # Convert video to 25fps temp_25fps_file = os.path.join(output_path, 'temp_25fps.mp4') convert_video_fps(input_file, temp_25fps_file, target_fps=25) # Detect faces filename_path = detectface(video_input_path=temp_25fps_file, output_path=output_path, detect_every_N_frame=detect_every_N_frame, scalar_face_detection=scalar_face_detection, number_of_speakers=number_of_speakers) # Extract audio audio_output = os.path.join(output_path, 'audio.wav') extract_audio(temp_25fps_file, audio_output, sample_rate=16000) # Crop mouth crop_mouth(video_direc=os.path.join(output_path, "faces"), landmark_direc=os.path.join(output_path, "landmark"), filename_path=filename_path, save_direc=os.path.join(output_path, "mouthroi"), convert_gray=True, testset_only=False) # Load model audiomodel = Dolphin.from_pretrained("JusperLee/Dolphin") audiomodel.cuda() audiomodel.eval() # Process each speaker with torch.no_grad(): for i in range(number_of_speakers): mouth_roi = np.load(os.path.join(output_path, "mouthroi", f"speaker{i+1}.npz"))["data"] mouth_roi = get_preprocessing_pipelines()["val"](mouth_roi) mix, sr = torchaudio.load(audio_output) mix = mix.cuda().mean(dim=0) window_size = 4 * sr hop_size = 4 * sr all_estimates = [] # 滑动窗口处理 start_idx = 0 while start_idx < len(mix): end_idx = min(start_idx + window_size, len(mix)) window_mix = mix[start_idx:end_idx] start_frame = int(start_idx / sr * 25) end_frame = int(end_idx / sr * 25) end_frame = min(end_frame, len(mouth_roi)) window_mouth_roi = mouth_roi[start_frame:end_frame] est_sources = audiomodel(window_mix[None], torch.from_numpy(window_mouth_roi[None, None]).float().cuda()) all_estimates.append({ 'start': start_idx, 'end': end_idx, 'estimate': est_sources[0].cpu() }) start_idx += hop_size if start_idx >= len(mix): break output_length = len(mix) merged_output = torch.zeros(1, output_length) weights = torch.zeros(output_length) for est in all_estimates: window_len = est['end'] - est['start'] hann_window = torch.hann_window(window_len) merged_output[0, est['start']:est['end']] += est['estimate'][0, :window_len] * hann_window weights[est['start']:est['end']] += hann_window merged_output[:, weights > 0] /= weights[weights > 0] torchaudio.save(os.path.join(output_path, f"speaker{i+1}_est.wav"), merged_output, sr) # Merge video with separated audio for each speaker output_files = [] for i in range(number_of_speakers): video_input = os.path.join(output_path, f"video_tracked{i+1}.mp4") audio_input = os.path.join(output_path, f"speaker{i+1}_est.wav") video_output = os.path.join(output_path, f"s{i+1}.mp4") merge_video_audio(video_input, audio_input, video_output) output_files.append(video_output) # Clean up temporary file if os.path.exists(temp_25fps_file): os.remove(temp_25fps_file) return output_files if __name__ == '__main__': parser = argparse.ArgumentParser(description='Video Speaker Separation using Dolphin model') parser.add_argument('--input', '-i', type=str, required=True, help='Path to input video file') parser.add_argument('--output', '-o', type=str, default=None, help='Output directory path (default: creates directory based on input filename)') parser.add_argument('--speakers', '-s', type=int, default=2, help='Number of speakers to separate (default: 2)') parser.add_argument('--detect-every-n', type=int, default=8, help='Detect faces every N frames (default: 8)') parser.add_argument('--face-scale', type=float, default=1.5, help='Face detection bounding box scale factor (default: 1.5)') parser.add_argument('--cuda-device', type=int, default=0, help='CUDA device ID to use (default: 0, set to -1 for CPU)') parser.add_argument('--config', type=str, default="checkpoints/vox2/conf.yml", help='Path to model configuration file') args = parser.parse_args() # 验证输入文件是否存在 if not os.path.exists(args.input): print(f"Error: Input file '{args.input}' does not exist") exit(1) # 如果没有指定输出路径，基于输入文件名创建输出目录 if args.output is None: input_basename = os.path.splitext(os.path.basename(args.input))[0] args.output = os.path.join(os.path.dirname(args.input), input_basename + "_output") # 设置CUDA设备 cuda_device = args.cuda_device if args.cuda_device >= 0 else None print(f"Processing video: {args.input}") print(f"Output directory: {args.output}") print(f"Number of speakers: {args.speakers}") print(f"CUDA device: {cuda_device if cuda_device is not None else 'CPU'}") # 处理视频 output_files = process_video( input_file=args.input, output_path=args.output, number_of_speakers=args.speakers, detect_every_N_frame=args.detect_every_n, scalar_face_detection=args.face_scale, config_path=args.config, cuda_device=cuda_device ) print("\nProcessing completed!") print("Output files:") for i, output_file in enumerate(output_files): print(f" Speaker {i+1}: {output_file}")