app.py

import gradio as gr
from llm_loader import load_model
from processing import process_input
from transcription_diarization import process_video
from visualization import create_charts
import os
import time
from config import hf_token, openai_api_key

# Load the model
llm = load_model(openai_api_key)

# Mapping of display names to language codes
LANGUAGE_MAP = {
    "English": "en",
    "Hebrew": "he",
    "Italian": "it",
    "French": "fr",
    "German": "de",
    "Chinese": "zh",
    "Arabic": "ar"
}


def analyze_video(video_path, language_display_name, max_speakers, progress=gr.Progress()):
    start_time = time.time()

    if not video_path:
        return "Please upload a video file.", gr.Textbox.update(value="Analysis not started.")

    # Convert the display name to the language code
    language = LANGUAGE_MAP[language_display_name]

    # Start the progress bar
    progress(0, desc="Starting analysis...")

    # Progress for diarization
    progress(0.2, desc="Starting diarization...")
    srt_path = process_video(video_path, hf_token, language, max_speakers)
    progress(0.4, desc="Diarization complete.")

    # Progress for transcription
    with open(srt_path, 'r', encoding='utf-8') as file:
        transcription = file.read()
    progress(0.6, desc="Transcription complete.")

    # Progress for processing the transcription
    progress(0.7, desc="Processing transcription...")
    results = process_input(transcription, llm)
    progress(0.8, desc="Transcription processing complete.")

    # Progress for creating charts
    progress(0.9, desc="Generating charts...")
    charts, explanations = create_charts(results)
    progress(1.0, desc="Charts generation complete.")

    # Clean up the temporary SRT file
    os.remove(srt_path)

    end_time = time.time()
    execution_time = end_time - start_time

    # Prepare outputs for each speaker
    output_components = []
    for speaker_id, speaker_charts in charts.items():
        speaker_explanations = explanations[speaker_id]

        output_components.extend([
            gr.Markdown(f"### Speaker {speaker_id}"),
            gr.Plot(speaker_charts["attachment"]),
            gr.Textbox(value=speaker_explanations["attachment"],
                       label=f"Attachment Styles Explanation - Speaker {speaker_id}", lines=2),
            gr.Plot(speaker_charts["dimensions"]),
            gr.Plot(speaker_charts["bigfive"]),
            gr.Textbox(value=speaker_explanations["bigfive"],
                       label=f"Big Five Traits Explanation - Speaker {speaker_id}", lines=2),
            gr.Plot(speaker_charts["personality"]),
            gr.Textbox(value=speaker_explanations["personality"],
                       label=f"Personality Disorders Explanation - Speaker {speaker_id}", lines=2)
        ])

    # Add the transcript and execution info at the end
    output_components.extend([
        gr.Textbox(value=transcription, label="Transcript", lines=10),
        gr.Textbox.update(value=f"Completed in {int(execution_time)} seconds.", label="Execution Information",
                          visible=True)
    ])

    return output_components, gr.Textbox.update(value=f"Completed in {int(execution_time)} seconds.")


# Define the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Video Analysis Tool")
    gr.Markdown("Upload a video to analyze speech patterns and personality traits.")

    video_input = gr.Video(label="Upload Video")
    language_input = gr.Dropdown(choices=list(LANGUAGE_MAP.keys()), value="English", label="Select Language")
    max_speakers = gr.Slider(minimum=1, maximum=4, step=1, value=2, label="Maximum Number of Speakers")

    analyze_button = gr.Button("Analyze")

    # Placeholder for dynamic outputs
    output_section = gr.Column()

    # Execution time box, initially displaying a waiting message
    execution_info_box = gr.Textbox(label="Execution Information", value="Waiting for analysis...", lines=2,
                                    visible=True)

    analyze_button.click(
        fn=analyze_video,
        inputs=[video_input, language_input, max_speakers],
        outputs=[output_section, execution_info_box],
        show_progress=True  # Enables the progress bar in Gradio
    )

# Launch the app
if __name__ == "__main__":
    iface.launch()