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, max_speakers, progress=gr.Progress()):
    start_time = time.time()
    if not video_path:
        return {"error": "Please upload a video file."}
    
    progress(0, desc="Starting analysis...")
    progress(0.2, desc="Starting transcription and diarization")
    transcription = diarize_audio(video_path, max_speakers)
    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 = create_charts(results)
    progress(1.0, desc="Charts generation complete.")
    
    end_time = time.time()
    execution_time = end_time - start_time
    
    output = {
        "transcript": transcription,
        "execution_info": f"Completed in {int(execution_time)} seconds.",
        "speakers": {}
    }
    
    for speaker_id, speaker_charts in charts.items():
        speaker_explanations = explanations[speaker_id]
        output["speakers"][speaker_id] = {
            "attachment": {
                "chart": speaker_charts["attachment"],
                "explanation": speaker_explanations["attachment"]
            },
            "dimensions": speaker_charts["dimensions"],
            "bigfive": {
                "chart": speaker_charts["bigfive"],
                "explanation": speaker_explanations["bigfive"]
            },
            "personality": {
                "chart": speaker_charts["personality"],
                "explanation": speaker_explanations["personality"]
            }
        }
    
    return output

def update_interface(result):
    if "error" in result:
        return [result["error"], gr.update(visible=False)] + [gr.update(visible=False)] * 21
    
    outputs = [
        result["transcript"],  # Transcript
        result["execution_info"]  # Execution info
    ]
    
    for i in range(3):
        speaker_id = f"Speaker {i+1}"
        if speaker_id in result["speakers"]:
            speaker_data = result["speakers"][speaker_id]
            outputs.extend([
                gr.update(value=speaker_data["attachment"]["chart"], visible=True),
                gr.update(value=speaker_data["attachment"]["explanation"], visible=True),
                gr.update(value=speaker_data["dimensions"], visible=True),
                gr.update(value=speaker_data["bigfive"]["chart"], visible=True),
                gr.update(value=speaker_data["bigfive"]["explanation"], visible=True),
                gr.update(value=speaker_data["personality"]["chart"], visible=True),
                gr.update(value=speaker_data["personality"]["explanation"], visible=True),
            ])
        else:
            outputs.extend([gr.update(visible=False)] * 7)
    
    return outputs

with gr.Blocks() as iface:
    gr.Markdown("# AI Personality Detection")
    gr.Markdown("Upload a video")
    
    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.Video(label="Upload Video")
            max_speakers = gr.Slider(minimum=1, maximum=3, step=1, value=2, label="Maximum Number of Speakers")
            analyze_button = gr.Button("Analyze")
        
        with gr.Column(scale=2):
            transcript_output = gr.Textbox(label="Transcript", lines=10)
            execution_info = gr.Textbox(label="Execution Information")
    
    with gr.Tabs() as tabs:
        speaker_tabs = [gr.Tab(f"Speaker {i+1}") for i in range(3)]
        speaker_outputs = []
        for tab in speaker_tabs:
            with tab:
                with gr.Column(scale=1):
                    attachment_plot = gr.Plot(label="Attachment Styles")
                    attachment_explanation = gr.Textbox(label="Attachment Styles Explanation")
                    dimensions_plot = gr.Plot(label="Attachment Dimensions")
                    bigfive_plot = gr.Plot(label="Big Five Traits")
                    bigfive_explanation = gr.Textbox(label="Big Five Traits Explanation")
                    personality_plot = gr.Plot(label="Personality Disorders")
                    personality_explanation = gr.Textbox(label="Personality Disorders Explanation")
                speaker_outputs.extend([
                    attachment_plot, attachment_explanation, dimensions_plot,
                    bigfive_plot, bigfive_explanation, personality_plot, personality_explanation
                ])

    def process_and_update(video_path, max_speakers):
        result = analyze_video(video_path, max_speakers)
        return update_interface(result)

    analyze_button.click(
        fn=process_and_update,
        inputs=[video_input, max_speakers],
        outputs=[transcript_output, execution_info] + speaker_outputs,
        show_progress=True
    )

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