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
import cv2
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 = []

    output_components.append(f"Completed in {int(execution_time)} seconds.")
    output_components.append(gr.Textbox(value=transcription, label="Transcript", lines=10, visible=True))

    for i, (speaker_id, speaker_charts) in enumerate(charts.items(), start=1):
        speaker_explanations = explanations[speaker_id]
        speaker_general_impression = general_impressions[speaker_id]
        
        output_components.extend([
            gr.Markdown(f"### {speaker_id}", visible=True),
            gr.Textbox(value=speaker_general_impression, label="General Impression", visible=True, lines=10),
            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, lines=2),
            gr.Plot(value=speaker_charts.get("bigfive", None), visible=True),
            gr.Textbox(value=speaker_explanations.get("bigfive", ""), label="Big Five Traits Explanation", visible=True, lines=2),
            gr.Plot(value=speaker_charts.get("personality", None), visible=True),
            gr.Textbox(value=speaker_explanations.get("personality", ""), label="Personality Disorders Explanation", visible=True, lines=2)
        ])

    # Pad with None for any missing speakers
    while len(output_components) < 29:
        output_components.append(None)

    return output_components

def use_example_1():
    return "examples/Scenes.From.A.Marriage.US.mp4"
def use_example_2():
    return "examples/Billie Eilish.mp4"
def use_example_3():
    return "examples/Elliot Rodger.mp4"

def get_middle_frame(video_path):
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    middle_frame_number = total_frames // 7
    cap.set(cv2.CAP_PROP_POS_FRAMES, middle_frame_number)
    ret, frame = cap.read()
    cap.release()
    if ret:
        preview_path = video_path.rsplit('.', 1)[0] + '_preview.jpg'
        cv2.imwrite(preview_path, frame)
        return preview_path
    return None

def clear_outputs():
    return [None] * 29

with gr.Blocks() as iface:
    gr.Markdown("# Multiple-Speakers-Personality-Analyzer")
    gr.Markdown("This project provides an advanced AI system designed for diagnosing and profiling personality attributes from video content based on a single speaker or multiple speakers in a conversation.")

    with gr.Row():
        video_input = gr.Video(label="Upload Video")
    
    analyze_button = gr.Button("Analyze")
            
    # Create output components
    output_components = []
    # Add transcript output near the top
    execution_box = gr.Textbox(label="Execution Info", value="N/A", lines=1)
    output_components.append(execution_box)

    transcript = gr.Textbox(label="Transcript", lines=10, visible=False)
    output_components.append(transcript)

    for n in range(3):  # Assuming maximum of 3 speakers
        with gr.Tab(label=f'Speaker {n + 1}', visible=False) as tab:
            gr.Markdown(visible=False)
            gr.Textbox(label="General Impression", visible=False)
            gr.Plot(visible=False)
            gr.Plot(visible=False)
            gr.Textbox(label="Attachment Styles Explanation", 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.extend([tab] + [component for component in tab.children])

    with open('description.txt', 'r') as file:
        description_txt = file.read()
    with gr.Tab(label=f'Description', visible=True):
        gr.Markdown(description_txt)
        gr.HTML("<div style='height: 20px;'></div>")
        gr.Image(value="appendix/AI Personality Detection flow - 1.png", label='Flowchart 1', width=1000)
        gr.Image(value="appendix/AI Personality Detection flow - 2.png", label='Flowchart 2', width=1000)

    gr.Markdown("### Example Videos")
    with gr.Row():
        with gr.Column(scale=1):
            example_video_1_path = "examples/Scenes.From.A.Marriage.US.mp4"
            preview_1 = get_middle_frame(example_video_1_path)
            gr.Image(preview_1, label="Scenes From A Marriage")
            example_video_1 = gr.Video(example_video_1_path, label="Example 1", visible=False)
            use_example_button_1 = gr.Button("Load Example 1")
        
        with gr.Column(scale=1):
            example_video_2_path = "examples/Billie Eilish.mp4"
            preview_2 = get_middle_frame(example_video_2_path)
            gr.Image(preview_2, label="Billie Eilish")
            example_video_2 = gr.Video(example_video_2_path, label="Example 2", visible=False)
            use_example_button_2 = gr.Button("Load Example 2")
        
        with gr.Column(scale=1):
            example_video_3_path = "examples/Elliot Rodger.mp4"
            preview_3 = get_middle_frame(example_video_3_path)
            gr.Image(preview_3, label="Elliot Rodger")
            example_video_3 = gr.Video(example_video_3_path, label="Example 3", visible=False)
            use_example_button_3 = gr.Button("Load Example 3")
            
    analyze_button.click(
        fn=clear_outputs,
        inputs=[],
        outputs=output_components
    ).then(
        fn=analyze_video,
        inputs=[video_input],
        outputs=output_components,
        show_progress=True
    )

    use_example_button_1.click(
        fn=use_example_1,
        inputs=[],
        outputs=[video_input],
    ).then(
        fn=clear_outputs,
        inputs=[],
        outputs=output_components
    ).then(
        fn=analyze_video,
        inputs=[video_input],
        outputs=output_components,
        show_progress=True
    )

    use_example_button_2.click(
        fn=use_example_2,
        inputs=[],
        outputs=[video_input],
    ).then(
        fn=clear_outputs,
        inputs=[],
        outputs=output_components
    ).then(
        fn=analyze_video,
        inputs=[video_input],
        outputs=output_components,
        show_progress=True
    )

    use_example_button_3.click(
        fn=use_example_3,
        inputs=[],
        outputs=[video_input],
    ).then(
        fn=clear_outputs,
        inputs=[],
        outputs=output_components
    ).then(
        fn=analyze_video,
        inputs=[video_input],
        outputs=output_components,
        show_progress=True
    )

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