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)
    print("Transcription:", transcription)  # Debug print
    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

    return {
        "transcript": transcription,
        "charts": charts,
        "explanations": explanations,
        "execution_time": int(execution_time)
    }

def create_output_components():
    with gr.Row() as row:
        with gr.Column():
            transcript = gr.Textbox(label="Transcript", lines=10, visible=False)
            
            tabs = gr.Tabs()
            for i in range(3):  # Pre-create 3 tabs (max number of speakers)
                with gr.Tab(f"Speaker {i+1}", visible=False) as tab:
                    gr.Markdown(f"## Speaker {i+1}", visible=False)
                    gr.Plot(label="Attachment", visible=False)
                    gr.Textbox(label="Attachment Styles Explanation", visible=False)
                    gr.Plot(label="Dimensions", visible=False)
                    gr.Plot(label="Big Five", visible=False)
                    gr.Textbox(label="Big Five Traits Explanation", visible=False)
                    gr.Plot(label="Personality", visible=False)
                    gr.Textbox(label="Personality Disorders Explanation", visible=False)
            
            execution_info = gr.Textbox(label="Execution Information", visible=True)
    return row, transcript, tabs, execution_info

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

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

    output_row, transcript_output, tabs_output, execution_info_output = create_output_components()

    def run_analysis(video_path, max_speakers):
        results = analyze_video(video_path, max_speakers)
        
        if "error" in results:
            return [
                "",  # transcript
                gr.Tabs(),  # tabs
                results["error"],  # execution_info
            ] + [gr.update(visible=False)] * 24  # 8 components per tab * 3 tabs
    
        transcript = results["transcript"]
        execution_info = f"Completed in {results['execution_time']} seconds."
    
        tab_updates = []
        for i in range(3):  # For each potential speaker
            if i < len(results["charts"]):
                speaker_id = list(results["charts"].keys())[i]
                speaker_charts = results["charts"][speaker_id]
                speaker_explanations = results["explanations"][speaker_id]
                
                tab_updates.extend([
                    gr.update(visible=True, label=speaker_id),  # Tab visibility and label
                    gr.update(value=f"## {speaker_id}", visible=True),  # Markdown
                    gr.update(value=speaker_charts.get("attachment", None), visible=True),  # Attachment plot
                    gr.update(value=speaker_explanations.get("attachment", ""), visible=True),  # Attachment explanation
                    gr.update(value=speaker_charts.get("dimensions", None), visible=True),  # Dimensions plot
                    gr.update(value=speaker_charts.get("bigfive", None), visible=True),  # Big Five plot
                    gr.update(value=speaker_explanations.get("bigfive", ""), visible=True),  # Big Five explanation
                    gr.update(value=speaker_charts.get("personality", None), visible=True),  # Personality plot
                    gr.update(value=speaker_explanations.get("personality", ""), visible=True),  # Personality explanation
                ])
            else:
                tab_updates.extend([gr.update(visible=False)] * 9)  # Hide unused tab and its components
    
        return [transcript, gr.Tabs.update(selected=0), execution_info] + tab_updates
    
    # Modify the analyze_button.click setup
    analyze_button.click(
        fn=run_analysis,
        inputs=[video_input, max_speakers],
        outputs=[
            transcript_output,
            tabs_output,
            execution_info_output
        ] + [component for tab in tabs_output.children for component in tab.children],
        show_progress=True
    )

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