app.py

# 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 [gr.Markdown("Please upload a video file.")] + [gr.update(visible=False)] * 48 + ["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_id}", visible=True),
            gr.update(value=speaker_charts.get("attachment", None), visible=True),
            gr.update(value=speaker_explanations.get("attachment", ""), visible=True),
            gr.update(value=speaker_charts.get("dimensions", None), visible=True),
            gr.update(value=speaker_charts.get("bigfive", None), visible=True),
            gr.update(value=speaker_explanations.get("bigfive", ""), visible=True),
            gr.update(value=speaker_charts.get("personality", None), visible=True),
            gr.update(value=speaker_explanations.get("personality", ""), visible=True),
        ])

    # Add the transcript at the end
    output_components.append(gr.update(value=transcription, visible=True))

    # Pad the output with invisible components if necessary
    while len(output_components) < 49:
        output_components.append(gr.update(visible=False))

    # Add the execution time
    output_components.append(f"Completed in {int(execution_time)} seconds.")

    return output_components

# Define the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# AI Personality Detection")
    gr.Markdown("Upload a video")

    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")

    # Create placeholders for output components
    output_components = []
    for _ in range(6):  # Assuming maximum of 6 speakers
        output_components.extend([
            gr.Markdown(visible=False),
            gr.Plot(visible=False),
            gr.Textbox(label="Attachment Styles Explanation", visible=False),
            gr.Plot(visible=False),
            gr.Plot(visible=False),
            gr.Textbox(label="Big Five Traits Explanation", visible=False),
            gr.Plot(visible=False),
            gr.Textbox(label="Personality Disorders Explanation", visible=False),
        ])
    output_components.append(gr.Textbox(label="Transcript", lines=10, visible=False))
    execution_info_box = gr.Textbox(label="Execution Information", value="Waiting for analysis...", lines=2)

    analyze_button.click(
        fn=analyze_video,
        inputs=[video_input, language_input, max_speakers],
        outputs=output_components + [execution_info_box],
        show_progress=True
    )

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