Inference_with_status.py

import warnings
warnings.filterwarnings("ignore")
import os
import torch
import torchaudio
import numpy as np
from moviepy import *
from PIL import Image, ImageDraw
import face_alignment
import cv2

from look2hear.models import Dolphin
from look2hear.datas.transform import get_preprocessing_pipelines

from face_detection_utils import detect_faces

# Import functions from original Inference.py
from Inference import (
    linear_interpolate, warp_img, apply_transform, cut_patch, convert_bgr2gray,
    save2npz, read_video, face2head, bb_intersection_over_union, 
    landmarks_interpolate, crop_patch, convert_video_fps, extract_audio, merge_video_audio
)

def detectface_with_status(video_input_path, output_path, detect_every_N_frame, scalar_face_detection, number_of_speakers, status_callback=None):
    """Face detection with status updates"""
    device = torch.device('cuda' if torch.cuda.get_device_name() else 'cpu')
    if status_callback:
        status_callback({'status': f'Running on device: {device}', 'progress': 0.0})
    
    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)
    if status_callback:
        status_callback({'status': f"Video: {video_clip.w}x{video_clip.h}, {video_clip.fps}fps", 'progress': 0.05})
    
    frames = [Image.fromarray(frame) for frame in video_clip.iter_frames()]
    total_frames = len(frames)
    if status_callback:
        status_callback({'status': f'Processing {total_frames} frames', 'progress': 0.1})
    
    video_clip.close()
    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False)
    
    for i, frame in enumerate(frames):
        if status_callback and i % 10 == 0:
            status_callback({'status': f'Tracking frame: {i+1}/{total_frames}', 'progress': 0.1 + 0.3 * (i / total_frames)})
        
        # 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 = np.asarray(detected_boxes, dtype=np.float32)
                areas = (detected_boxes[:, 2] - detected_boxes[:, 0]) * (detected_boxes[:, 3] - detected_boxes[:, 1])
                sort_idx = np.argsort(areas)[::-1]
                detected_boxes = detected_boxes[sort_idx][:number_of_speakers]
                detected_boxes = face2head(detected_boxes, scalar_face_detection)
                detected_boxes = [box for box in detected_boxes]
            else:
                detected_boxes = []

        # Process the detection results (same as original)
        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)]
    if status_callback:
        status_callback({'status': f"Frame counts per speaker: {frame_counts}", 'progress': 0.4})
    
    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):
        if status_callback:
            status_callback({'status': f'Saving tracked video for speaker {s+1}', 'progress': 0.4 + 0.1 * (s / 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):
        # Create landmark directory if it doesn't exist
        landmark_dir = os.path.join(output_path, 'landmark')
        os.makedirs(landmark_dir, exist_ok=True)
        save2npz(os.path.join(landmark_dir, '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]
    filename_dir = os.path.join(output_path, 'filename_input')
    os.makedirs(filename_dir, exist_ok=True)
    csvfile = open(os.path.join(filename_dir, 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(filename_dir, str(video_name) + '.csv')


def crop_mouth_with_status(video_direc, landmark_direc, filename_path, save_direc, status_callback=None, convert_gray=False, testset_only=False):
    """Crop mouth with status updates"""
    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(',')
        
        if status_callback:
            status_callback({'status': f'Processing speaker{int(person_id)+1}', 'progress': 0.5 + 0.1 * filename_idx / len(lines)})

        video_pathname = os.path.join(video_direc, filename+'.mp4')
        landmarks_pathname = os.path.join(landmark_direc, filename+'.npz')
        
        # Create mouthroi directory if it doesn't exist
        os.makedirs(save_direc, exist_ok=True)
        dst_pathname = os.path.join(save_direc, filename+'.npz')

        multi_sub_landmarks = np.load(landmarks_pathname, allow_pickle=True)['data']
        if len(multi_sub_landmarks) == 0:
            print(f"No landmarks found for {filename}, skipping crop.")
            continue

        landmark_frame_count = len(multi_sub_landmarks)
        cap = cv2.VideoCapture(video_pathname)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
        cap.release()

        if frame_count > 0 and frame_count != landmark_frame_count:
            print(
                f"Frame count mismatch for {filename}: video has {frame_count} frames, "
                f"landmarks have {landmark_frame_count} entries. Adjusting to match."
            )
            if frame_count < landmark_frame_count:
                multi_sub_landmarks = multi_sub_landmarks[:frame_count]
            else:
                pad_count = frame_count - landmark_frame_count
                pad = np.repeat(multi_sub_landmarks[-1:], pad_count, axis=0)
                multi_sub_landmarks = np.concatenate((multi_sub_landmarks, pad), axis=0)

        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)]
            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 process_video_with_status(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, status_callback=None):
    """Main processing function with status updates"""
    
    # 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
    if status_callback:
        status_callback({'status': 'Converting video to 25fps', 'progress': 0.0})
    
    temp_25fps_file = os.path.join(output_path, 'temp_25fps.mp4')
    convert_video_fps(input_file, temp_25fps_file, target_fps=25)
    
    # Detect faces
    if status_callback:
        status_callback({'status': 'Detecting faces and tracking speakers', 'progress': 0.1})
    
    filename_path = detectface_with_status(
        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,
        status_callback=status_callback
    )
    torch.cuda.empty_cache()
    # Extract audio
    if status_callback:
        status_callback({'status': 'Extracting audio from video', 'progress': 0.5})
    
    audio_output = os.path.join(output_path, 'audio.wav')
    extract_audio(temp_25fps_file, audio_output, sample_rate=16000)
    
    # Crop mouth
    if status_callback:
        status_callback({'status': 'Cropping mouth regions', 'progress': 0.55})
    
    crop_mouth_with_status(
        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,
        status_callback=status_callback
    )
    
    # Load model
    if status_callback:
        status_callback({'status': 'Loading Dolphin model', 'progress': 0.6})
    torch.cuda.empty_cache()
    audiomodel = Dolphin.from_pretrained("JusperLee/Dolphin")
    # audiomodel.cuda()
    audiomodel.eval()
    
    # Process each speaker
    with torch.no_grad():
        for i in range(number_of_speakers):
            if status_callback:
                status_callback({'status': f'Processing audio for speaker {i+1}', 'progress': 0.65 + 0.25 * (i / number_of_speakers)})
            
            mouth_roi_path = os.path.join(output_path, "mouthroi", f"speaker{i+1}.npz")
            mouth_roi = np.load(mouth_roi_path)["data"]
            mouth_roi = get_preprocessing_pipelines()["val"](mouth_roi)
            
            mix, sr = torchaudio.load(audio_output)
            mix = mix.mean(dim=0)
            
            window_size = 4 * sr 
            hop_size = int(4 * sr)
            
            all_estimates = []
            
            # Sliding window processing
            start_idx = 0
            window_count = 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())
                
                all_estimates.append({
                    'start': start_idx,
                    'end': end_idx,
                    'estimate': est_sources[0].cpu()
                })
                
                window_count += 1
                if status_callback:
                    progress = 0.65 + 0.25 * (i / number_of_speakers) + 0.25 / number_of_speakers * (window_count * hop_size / len(mix))
                    status_callback({'status': f'Processing audio window {window_count} for speaker {i+1}', 'progress': min(progress, 0.9)})
                
                start_idx += hop_size
                
                if start_idx >= len(mix):
                    break
                torch.cuda.empty_cache()
            
            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]
            
            audio_save_path = os.path.join(output_path, f"speaker{i+1}_est.wav")
            torchaudio.save(audio_save_path, merged_output, sr)

    # Merge video with separated audio for each speaker
    torch.cuda.empty_cache()
    if status_callback:
        status_callback({'status': 'Merging videos with separated audio', 'progress': 0.9})
    
    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)
    
    if status_callback:
        status_callback({'status': 'Processing completed!', 'progress': 1.0})
    
    return output_files