app.py

import gradio as gr
from llm_loader import load_model
from processing import process_input
from transcription_diarization import diarize_audio
from visualization import create_charts
import time
from config import openai_api_key

# Load the model
llm = load_model(openai_api_key)

def analyze_video(video_path, progress=gr.Progress()):
    start_time = time.time()
    if not video_path:
        return [None] * 29  # Return None for all outputs

    progress(0, desc="Starting analysis...")
    progress(0.2, desc="Starting transcription and diarization")
    transcription = diarize_audio(video_path)
    progress(0.5, desc="Transcription and diarization complete.")

    progress(0.6, desc="Processing transcription")
    results = process_input(transcription, llm)
    progress(0.7, desc="Transcription processing complete.")

    progress(0.9, desc="Generating charts")
    charts, explanations, general_impressions = create_charts(results)
    progress(1.0, desc="Charts generation complete.")

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

    output_components = [transcription]  # transcript

    for i, (speaker_id, speaker_charts) in enumerate(charts.items(), start=1):
        speaker_explanations = explanations[speaker_id]
        speaker_general_impression = general_impressions[speaker_id]

        speaker_section = [
            gr.Markdown(f"## {speaker_id}", visible=True),
            speaker_general_impression,  # speaker impression
            gr.Textbox(value=speaker_general_impression.get("general_impression", ""), label="General Impression",
                       visible=True),
            gr.Plot(value=speaker_charts.get("attachment", None), visible=True),
            gr.Plot(value=speaker_charts.get("dimensions", None), visible=True),

            gr.Textbox(value=speaker_explanations.get("attachment", ""), label="Attachment Styles Explanation",
                       visible=True),
            gr.Plot(value=speaker_charts.get("bigfive", None), visible=True),
            gr.Textbox(value=speaker_explanations.get("bigfive", ""), label="Big Five Traits Explanation",
                       visible=True),
            gr.Plot(value=speaker_charts.get("personality", None), visible=True),
            gr.Textbox(value=speaker_explanations.get("personality", ""), label="Personality Disorders Explanation",
                       visible=True),
        ]
        output_components.extend(speaker_section)
        
    # Pad with None for any missing speakers
    while len(output_components) < 28:
        output_components.extend([gr.update(visible=False)] * 9)

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

    return output_components

def update_output(*args):
    return [gr.update(value=arg, visible=arg is not None) for arg in args]

def use_example():
    return "examples/Scenes.From.A.Marriage.US.mp4"

with gr.Blocks() as iface:
    gr.Markdown("# AI Personality Detection")
    
    with gr.Row():
        with gr.Column(scale=3):
            gr.Markdown("Upload a video")
            video_input = gr.Video(label="Upload Video")
            analyze_button = gr.Button("Analyze")
        with gr.Column(scale=1):
            gr.Markdown("Example Video")
            example_video = gr.Video("examples/Scenes.From.A.Marriage.US.mp4", label="Example Video")
            use_example_button = gr.Button("Use Example Video")

    # Create placeholder components for output
    with gr.Column() as output_container:
        transcript_output = gr.Textbox(label="Transcript", lines=10, visible=False)
        speaker_outputs = []
        for i in range(1, 4):
            with gr.Column(visible=False) as speaker_column:
                speaker_header = gr.Markdown(f"## Speaker {i}")
                speaker_impression = gr.Textbox(label="General Impression", lines=3)
                speaker_attachment = gr.Plot(label="Attachment Styles")
                speaker_attachment_exp = gr.Textbox(label="Attachment Styles Explanation")
                speaker_dimensions = gr.Plot(label="Attachment Dimensions")
                speaker_bigfive = gr.Plot(label="Big Five Traits")
                speaker_bigfive_exp = gr.Textbox(label="Big Five Traits Explanation")
                speaker_personality = gr.Plot(label="Personality Disorders")
                speaker_personality_exp = gr.Textbox(label="Personality Disorders Explanation")
            speaker_outputs.extend([
                speaker_header, speaker_impression, speaker_attachment, speaker_attachment_exp,
                speaker_dimensions, speaker_bigfive, speaker_bigfive_exp, speaker_personality,
                speaker_personality_exp
            ])
        execution_info = gr.Textbox(label="Execution Information", visible=True)

    all_outputs = [transcript_output] + speaker_outputs + [execution_info]

    analyze_button.click(
        fn=analyze_video,
        inputs=[video_input],
        outputs=all_outputs,
        show_progress=True
    ).then(
        fn=update_output,
        inputs=all_outputs,
        outputs=all_outputs
    )

    use_example_button.click(
        fn=use_example,
        inputs=[],
        outputs=[video_input],
    ).then(
        fn=analyze_video,
        inputs=[video_input],
        outputs=all_outputs,
        show_progress=True
    ).then(
        fn=update_output,
        inputs=all_outputs,
        outputs=all_outputs
    )

if __name__ == "__main__":
    iface.launch()