app.py

import gradio as gr
import pandas as pd
import cv2
import torch
import tempfile
import os
import librosa
from fer import FER
from transformers import AutoModelForAudioClassification, pipeline
from moviepy.editor import VideoFileClip, AudioFileClip
import numpy as np
from torch.nn.functional import softmax
import whisper_timestamped as whisper
from translate import Translator

# Load pre-trained models
audio_model = AutoModelForAudioClassification.from_pretrained("3loi/SER-Odyssey-Baseline-WavLM-Categorical-Attributes", trust_remote_code=True)
face_detector = FER(mtcnn=True)
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None)

# Set mean and std for audio model
mean = audio_model.config.mean
std = audio_model.config.std

# Function to extract audio from video for audio emotion analysis
def extract_audio_from_video(video_path):
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file:
        video_clip = VideoFileClip(video_path)
        audio_clip = video_clip.audio
        audio_clip.write_audiofile(temp_audio_file.name, codec="pcm_s16le")
        return temp_audio_file.name

# Function to perform audio emotion detection per second
def process_audio_and_detect_emotions(audio_clip):
    audio_np = np.array(audio_clip)
    mask = torch.ones(1, len(audio_np))
    wavs = torch.tensor(audio_np).unsqueeze(0)

    with torch.no_grad():
        pred = audio_model(wavs, mask)
    logits = pred.logits if hasattr(pred, 'logits') else pred[0]
    labels = {0: 'Angry', 1: 'Sad', 2: 'Happy', 3: 'Surprise', 4: 'Fear', 5: 'Disgust', 7: 'Neutral'}
    probabilities = softmax(logits, dim=-1).squeeze(0)[[0, 1, 2, 3, 4, 5, 7]]
    probabilities = probabilities / probabilities.sum()
    df = pd.DataFrame([probabilities.numpy()], columns=labels.values())
    return df

# Function to analyze audio emotions
def analyze_audio_emotions(video_path):
    temp_audio_path = None
    try:
        temp_audio_path = extract_audio_from_video(video_path)
        raw_wav, _ = librosa.load(temp_audio_path, sr=audio_model.config.sampling_rate)
        norm_wav = (raw_wav - mean) / (std + 0.000001)

        times = []
        emotions_dfs = []
        for start_time in range(0, len(norm_wav), audio_model.config.sampling_rate):
            audio_segment = norm_wav[start_time:start_time + audio_model.config.sampling_rate]
            df = process_audio_and_detect_emotions(audio_segment)
            times.append(start_time / audio_model.config.sampling_rate)
            emotions_dfs.append(df)

        emotions_df = pd.concat(emotions_dfs, ignore_index=True)
        emotions_df.insert(0, "Time(s)", times)
        emotion_rename_map = {'Angry': 'anger', 'Sad': 'sadness', 'Happy': 'happy', 'Surprise': 'surprise', 'Fear': 'fear', 'Disgust': 'disgust', 'Neutral': 'neutral'}
        emotions_df.rename(columns=emotion_rename_map, inplace=True)

        emotions_xlsx_path = tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False).name
        emotions_df.to_excel(emotions_xlsx_path, index=False)

        return f"Audio emotion detection completed successfully.", emotions_df, emotions_xlsx_path

    except Exception as e:
        return f"Error during audio emotion detection: {str(e)}", None, None
    finally:
        if temp_audio_path and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)

# Function to detect facial emotions
def detect_faces_and_emotions(video_path):
    temp_video_path = None
    temp_audio_path = None
    output_video_path = None
    emotions_data = []
    try:
        temp_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        temp_video_path = temp_video.name
        temp_audio = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
        temp_audio_path = temp_audio.name
        output_xlsx = tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False)
        output_xlsx_path = output_xlsx.name

        original_video = VideoFileClip(video_path)
        original_audio = original_video.audio
        original_audio.write_audiofile(temp_audio_path)

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise Exception("Error: Could not open video file.")

        fps = int(cap.get(cv2.CAP_PROP_FPS))
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(temp_video_path, fourcc, fps, (frame_width, frame_height))

        frame_number = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            if frame is None:
                continue

            time_seconds = round(frame_number / fps)
            result = face_detector.detect_emotions(frame)

            for face in result:
                bounding_box = face["box"]
                emotions = face["emotions"]
                emotions["Time(s)"] = time_seconds
                emotions_data.append(emotions)
                cv2.rectangle(frame, (bounding_box[0], bounding_box[1]),
                              (bounding_box[0] + bounding_box[2], bounding_box[1] + bounding_box[3]), (0, 155, 255), 2)
                for index, (emotion_name, score) in enumerate(emotions.items()):
                    color = (211, 211, 211) if score < 0.01 else (255, 0, 0)
                    emotion_score = "{}: {:.2f}".format(emotion_name, score)
                    cv2.putText(frame, emotion_score, (bounding_box[0], bounding_box[1] + bounding_box[3] + 30 + index * 15),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)

            out.write(frame)
            frame_number += 1

        cap.release()
        out.release()

        emotions_df = pd.DataFrame(emotions_data)
        emotions_df['Time(s)'] = emotions_df['Time(s)'].round().astype(int)
        max_time = emotions_df['Time(s)'].max()
        all_times = pd.DataFrame({'Time(s)': range(max_time + 1)})
        avg_scores = emotions_df.groupby("Time(s)").mean().reset_index()
        df_merged = pd.merge(all_times, avg_scores, on='Time(s)', how='left')
        df_merged.fillna(0, inplace=True)
        df_merged['Time(s)'] = df_merged['Time(s)'].astype(str) + " sec"
        df_merged.to_excel(output_xlsx_path, index=False)

        processed_video = VideoFileClip(temp_video_path)
        audio = AudioFileClip(temp_audio_path)
        final_video = processed_video.set_audio(audio)
        output_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        output_video_path = output_video.name
        final_video.write_videofile(output_video_path, codec='libx264')

        return "Face and emotion detection completed successfully.", df_merged, output_xlsx_path, output_video_path

    except Exception as e:
        return f"Error during processing: {str(e)}", None, None, None
    finally:
        if temp_video_path and os.path.exists(temp_video_path):
            os.remove(temp_video_path)
        if temp_audio_path and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)

# Function to analyze text emotions
def process_video_text(video_path):
    temp_audio_path = None
    try:
        video_clip = VideoFileClip(video_path)
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file:
            temp_audio_path = temp_audio_file.name
            video_clip.audio.write_audiofile(temp_audio_path)

        audio = whisper.load_audio(temp_audio_path)
        model = whisper.load_model("medium", device="cpu")
        result = whisper.transcribe(model, audio)

        # Create lists to store word-level data with timestamps
        word_texts = []
        word_starts = []
        word_ends = []
        word_confidences = []

        for segment in result['segments']:
            for word in segment['words']:
                word_texts.append(word['text'])
                word_starts.append(word['start'])
                word_ends.append(word['end'])
                word_confidences.append(word['confidence'])

        # Create segments DataFrame
        segments_data = [{'text': seg['text'], 'start': seg['start'], 'end': seg['end'], 'confidence': seg['confidence']} for seg in result['segments']]
        segments_df = pd.DataFrame(segments_data)

        # Translate from Korean to English
        translator = Translator(from_lang='ko', to_lang='en')
        segments_df['Translated_Text'] = segments_df['text'].apply(lambda x: translator.translate(x))

        # Apply the sentiment analysis model to the translated text
        segments_df['Sentiment_Scores'] = segments_df['Translated_Text'].apply(lambda x: {entry['label']: entry['score'] for entry in classifier(x)[0]})

        # Split the sentiment scores into individual columns
        sentiment_df = segments_df['Sentiment_Scores'].apply(pd.Series)
        sentiment_df = pd.concat([segments_df, sentiment_df], axis=1)

        # Create words DataFrame
        words_data = {
            'text': word_texts,
            'start': word_starts,
            'end': word_ends,
            'confidence': word_confidences
        }
        words_df = pd.DataFrame(words_data)

        # Round up the start time to the next second
        words_df['second'] = words_df['start'].apply(lambda x: int(np.ceil(x)))

        # Group words by second, concatenating words that belong to the same second
        words_grouped = words_df.groupby('second').agg({
            'text': lambda x: ' '.join(x),
            'start': 'min',
            'end': 'max',
            'confidence': 'mean'
        }).reset_index()

        # Fill in missing seconds
        max_second = int(video_clip.duration)  # The last second in the video
        all_seconds = pd.DataFrame({'second': np.arange(0, max_second + 1)})  # Start from 0 and go to the maximum second
        words_grouped = all_seconds.merge(words_grouped, on='second', how='left')

        # Fill missing values with blanks or zeros
        words_grouped['text'].fillna('', inplace=True)
        words_grouped.fillna(0, inplace=True)

        # Initialize emotion columns with NaN values
        emotion_columns = sentiment_df.columns.difference(['text', 'start', 'end', 'confidence', 'Translated_Text', 'Sentiment_Scores'])
        for col in emotion_columns:
            words_grouped[col] = np.nan

        # For each second, find the corresponding segment and copy its emotion scores
        for i, row in words_grouped.iterrows():
            matching_segment = sentiment_df[(sentiment_df['start'] <= row['start']) & (sentiment_df['end'] >= row['end'])]
            if not matching_segment.empty:
                for emotion in emotion_columns:
                    words_grouped.at[i, emotion] = matching_segment.iloc[0][emotion]

        # Replace any NaN values in emotion columns with 0
        words_grouped[emotion_columns] = words_grouped[emotion_columns].fillna(0)

        # Save DataFrames to XLSX files
        segments_xlsx_path = tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False).name
        words_xlsx_path = tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False).name
        sentiment_df.to_excel(segments_xlsx_path, index=False)
        words_grouped.to_excel(words_xlsx_path, index=False)

        return words_grouped, sentiment_df, words_xlsx_path, segments_xlsx_path, "Text emotion processing completed successfully!"

    except Exception as e:
        return None, None, None, None, f"Error during text emotion processing: {str(e)}"
    finally:
        if temp_audio_path and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)

# Gradio App
def gradio_app():
    interface = gr.Blocks()

    with interface:
        gr.Markdown("## I-MEQ: Emotion Monitoring System")
        video_input = gr.Video(label="Upload your video for analysis", height=600)

        with gr.Row():
            analyze_audio_button = gr.Button("Analyze Audio Emotions")
            analyze_fer_button = gr.Button("Analyze Facial Emotions")
            analyze_text_button = gr.Button("Transcribe & Analyze Textual Emotions")

        with gr.Row():
            with gr.Column():
                audio_analysis_status = gr.Textbox(label="Audio Emotion Analysis Status")
                audio_emotions_dataframe = gr.Dataframe(label="Audio Emotions DataFrame", interactive=False)
                audio_emotions_xlsx_download = gr.File(label="Download Audio Emotions XLSX")

            with gr.Column():
                fer_analysis_status = gr.Textbox(label="Facial Emotion Analysis Status")
                fer_emotions_dataframe = gr.Dataframe(label="Facial Emotions DataFrame", interactive=False)
                fer_emotions_xlsx_download = gr.File(label="Download Facial Emotions XLSX")
                processed_video_download = gr.File(label="Download Processed Video")

            with gr.Column():
                text_analysis_status = gr.Textbox(label="Text Sentiment Analysis Status")
                words_dataframe = gr.Dataframe(label="Words DataFrame", interactive=False)
                segments_dataframe = gr.Dataframe(label="Segments DataFrame", interactive=False)
                words_xlsx_download = gr.File(label="Download Words XLSX")
                segments_xlsx_download = gr.File(label="Download Segments XLSX")

        analyze_audio_button.click(
            analyze_audio_emotions,
            inputs=video_input,
            outputs=[
                audio_analysis_status,
                audio_emotions_dataframe,
                audio_emotions_xlsx_download
            ]
        )

        analyze_fer_button.click(
            detect_faces_and_emotions,
            inputs=video_input,
            outputs=[
                fer_analysis_status,
                fer_emotions_dataframe,
                fer_emotions_xlsx_download,
                processed_video_download
            ]
        )

        analyze_text_button.click(
            process_video_text,
            inputs=video_input,
            outputs=[
                words_dataframe,
                segments_dataframe,
                words_xlsx_download,
                segments_xlsx_download,
                text_analysis_status
            ]
        )

    interface.launch()

# Start the Gradio app
gradio_app()