app.py

import os
import subprocess
import gradio as gr
import whisper
import yt_dlp
import torch
import numpy as np
from moviepy.editor import VideoFileClip
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import BlipProcessor, BlipForConditionalGeneration
import cv2

# Define the necessary functions

def download_youtube_video(video_url, output_path):
    ydl_opts = {
        'format': 'bestvideo+bestaudio',
        'outtmpl': os.path.join(output_path, '%(title)s.%(ext)s'),
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([video_url])
        video_info = ydl.extract_info(video_url, download=False)
        video_title = video_info.get('title', 'video')
        return os.path.join(output_path, f"{video_title}.webm")

def convert_to_mp4(input_path, output_path):
    output_file = os.path.join(output_path, 'video.mp4')
    command = ['ffmpeg', '-i', input_path, '-c', 'copy', output_file]
    subprocess.run(command, check=True)
    return output_file

def extract_audio_from_video(video_path):
    video_clip = VideoFileClip(video_path)
    audio_output = os.path.join(output_path, 'audio.mp3')
    audio_clip = video_clip.audio
    audio_clip.write_audiofile(audio_output)
    return audio_output

def convert_mp3_to_wav(mp3_path):
    from pydub import AudioSegment
    audio = AudioSegment.from_mp3(mp3_path)
    wav_output = os.path.join(output_path, 'audio.wav')
    audio.export(wav_output, format="wav")
    return wav_output

def process_text(text):
    model_name = "cardiffnlp/twitter-roberta-base-emotion"
    emotion_labels = ['anger', 'joy', 'optimism', 'sad']
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)

    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    emotion_probs = torch.softmax(logits, dim=-1).squeeze()
    predicted_emotion = emotion_labels[torch.argmax(emotion_probs)]
    
    emotion_dict = {emotion_labels[i]: emotion_probs[i].item() for i in range(len(emotion_labels))}
    
    return emotion_dict, predicted_emotion

def preprocess_frame(frame):
    frame = cv2.resize(frame, (224, 224))
    pixel_values = caption_processor(images=frame, return_tensors="pt").pixel_values
    return pixel_values

def generate_caption(pixel_values):
    caption_ids = caption_model.generate(pixel_values)
    caption = caption_processor.batch_decode(caption_ids, skip_special_tokens=True)[0]
    return caption

def predict_emotions(caption):
    inputs = emotion_tokenizer(caption, return_tensors='pt', truncation=True, padding=True)
    outputs = emotion_model(**inputs)
    
    emotion_probs = torch.softmax(outputs.logits, dim=1)
    
    predicted_emotions = {label: prob.item() for label, prob in zip(emotion_labels, emotion_probs[0])}
    
    return predicted_emotions

# Load models and processors once at the start
caption_model_name = "Salesforce/blip-image-captioning-base"
caption_processor = BlipProcessor.from_pretrained(caption_model_name)
caption_model = BlipForConditionalGeneration.from_pretrained(caption_model_name)

emotion_model_name = "j-hartmann/emotion-english-distilroberta-base"
emotion_tokenizer = AutoTokenizer.from_pretrained(emotion_model_name)
emotion_model = AutoModelForSequenceClassification.from_pretrained(emotion_model_name)

# Gradio Interface Function
def analyze_video(video_url):
    # Set output path for downloads
    global output_path
    output_path = './'

    # Download the video
    video_path = download_youtube_video(video_url, output_path)

    # Convert to mp4 format
    mp4_path = convert_to_mp4(video_path, output_path)

    # Extract audio from the video
    audio_path = extract_audio_from_video(mp4_path)

    # Convert audio to wav format for processing
    audio_wav_path = convert_mp3_to_wav(audio_path)

    # Process the audio using Whisper for transcription
    model_whisper = whisper.load_model("base")
    
    result_whisper = model_whisper.transcribe(audio_wav_path)
    
    transcript = result_whisper['text']
    
    # Process text to get emotions
    emotion_dict_text, predicted_emotion_text = process_text(transcript)

    
   # Process the video using image captioning and emotion recognition
   n_frame_interval = 60  # Process every 60th frame
   emotion_vectors_video = []

   # Process the video frames for emotions using BLIP model 
   video_capture = cv2.VideoCapture(mp4_path)
   total_frames_video = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

   frame_count_video = 0

   while video_capture.isOpened():
       ret_video, frame_video = video_capture.read()

       if not ret_video or frame_count_video > total_frames_video:
           break

       if frame_count_video % n_frame_interval == 0:
           pixel_values_video = preprocess_frame(frame_video)
           caption_video = generate_caption(pixel_values_video)
           predicted_emotions_video, _ = predict_emotions(caption_video)

           # Collect emotion vectors from frames 
           emotion_vectors_video.append(np.array(list(predicted_emotions_video.values())))

       frame_count_video += 1

   video_capture.release()

   # Aggregate results from video frames 
   average_emotion_vector_video = np.mean(emotion_vectors_video, axis=0)

   # Combine text and video emotion results 
   combined_emotion_vector_final= np.concatenate((np.array(list(emotion_dict_text.values())), average_emotion_vector_video))

   final_most_predicted_index= np.argmax(combined_emotion_vector_final)
   
   final_most_predicted_emotion= list(emotion_dict_text.keys())[final_most_predicted_index]

   return transcript, predicted_emotion_text, final_most_predicted_emotion


# Create Gradio interface 
iface= gr.Interface(fn=analyze_video,
                     inputs=gr.Textbox(label="YouTube Video URL"),
                     outputs=["text", "text", "text"],
                     title="Multimodal Emotion Recognition",
                     description="Enter a YouTube Video URL to analyze emotions from both audio and visual content.")
                     
# Launch the app 
if __name__ == "__main__":
     iface.launch()