visualization.py

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.colors as mcolors
from matplotlib.colors import LinearSegmentedColormap
import seaborn as sns
import numpy as np
import pandas as pd
import cv2
from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip, VideoClip, concatenate_videoclips
from moviepy.video.fx.all import resize
from PIL import Image, ImageDraw, ImageFont
from matplotlib.patches import Rectangle
from scipy import interpolate
import os

# Utility functions
def seconds_to_timecode(seconds):
    hours = seconds // 3600
    minutes = (seconds % 3600) // 60
    seconds = seconds % 60
    return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}"

def determine_anomalies(values, threshold):
    mean = np.mean(values)
    std = np.std(values)
    anomalies = np.where(values > mean + threshold * std)[0]
    return anomalies

def plot_mse(df, mse_values, title, color='navy', time_threshold=3, anomaly_threshold=4):
    plt.figure(figsize=(16, 8), dpi=300)
    fig, ax = plt.subplots(figsize=(16, 8))

    if 'Seconds' not in df.columns:
        df['Seconds'] = df['Timecode'].apply(
            lambda x: sum(float(t) * 60 ** i for i, t in enumerate(reversed(x.split(':')))))

    # Ensure df and mse_values have the same length and remove NaN values
    min_length = min(len(df), len(mse_values))
    df = df.iloc[:min_length].copy()
    mse_values = mse_values[:min_length]

    # Remove NaN values and create a mask for valid data
    valid_mask = ~np.isnan(mse_values)
    df = df[valid_mask]
    mse_values = mse_values[valid_mask]

    # Function to identify continuous segments
    def get_continuous_segments(seconds, values, max_gap=1):
        segments = []
        current_segment = []
        for i, (sec, val) in enumerate(zip(seconds, values)):
            if not current_segment or (sec - current_segment[-1][0] <= max_gap):
                current_segment.append((sec, val))
            else:
                segments.append(current_segment)
                current_segment = [(sec, val)]
        if current_segment:
            segments.append(current_segment)
        return segments

    # Get continuous segments
    segments = get_continuous_segments(df['Seconds'], mse_values)

    # Plot each segment separately
    for segment in segments:
        segment_seconds, segment_mse = zip(*segment)
        ax.scatter(segment_seconds, segment_mse, color=color, alpha=0.3, s=5)
        
        # Calculate and plot rolling mean and std for this segment
        if len(segment) > 1:  # Only if there's more than one point in the segment
            segment_df = pd.DataFrame({'Seconds': segment_seconds, 'MSE': segment_mse})
            segment_df = segment_df.sort_values('Seconds')
            mean = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).mean()
            std = segment_df['MSE'].rolling(window=min(10, len(segment)), min_periods=1, center=True).std()
            
            ax.plot(segment_df['Seconds'], mean, color=color, linewidth=0.5)
            ax.fill_between(segment_df['Seconds'], mean - std, mean + std, color=color, alpha=0.1)

    median = np.median(mse_values)
    ax.axhline(y=median, color='black', linestyle='--', label='Median Baseline')

    threshold = np.mean(mse_values) + anomaly_threshold * np.std(mse_values)
    ax.axhline(y=threshold, color='red', linestyle='--', label=f'Anomaly Threshold')
    ax.text(ax.get_xlim()[1], threshold, f'Anomaly Threshold', verticalalignment='center', horizontalalignment='left', color='red')

    anomalies = determine_anomalies(mse_values, anomaly_threshold)
    anomaly_frames = df['Frame'].iloc[anomalies].tolist()

    ax.scatter(df['Seconds'].iloc[anomalies], mse_values[anomalies], color='red', s=20, zorder=5)

    anomaly_data = list(zip(df['Timecode'].iloc[anomalies],
                            df['Seconds'].iloc[anomalies],
                            mse_values[anomalies]))
    anomaly_data.sort(key=lambda x: x[1])

    grouped_anomalies = []
    current_group = []
    for timecode, sec, mse in anomaly_data:
        if not current_group or sec - current_group[-1][1] <= time_threshold:
            current_group.append((timecode, sec, mse))
        else:
            grouped_anomalies.append(current_group)
            current_group = [(timecode, sec, mse)]
    if current_group:
        grouped_anomalies.append(current_group)

    for group in grouped_anomalies:
        start_sec = group[0][1]
        end_sec = group[-1][1]
        rect = Rectangle((start_sec, ax.get_ylim()[0]), end_sec - start_sec, ax.get_ylim()[1] - ax.get_ylim()[0],
                         facecolor='red', alpha=0.2, zorder=1)
        ax.add_patch(rect)

    for group in grouped_anomalies:
        highest_mse_anomaly = max(group, key=lambda x: x[2])
        timecode, sec, mse = highest_mse_anomaly
        ax.annotate(timecode, (sec, mse), textcoords="offset points", xytext=(0, 10),
                    ha='center', fontsize=6, color='red')

    max_seconds = df['Seconds'].max()
    num_ticks = 100
    tick_locations = np.linspace(0, max_seconds, num_ticks)
    tick_labels = [seconds_to_timecode(int(s)) for s in tick_locations]

    ax.set_xticks(tick_locations)
    ax.set_xticklabels(tick_labels, rotation=90, ha='center', fontsize=6)

    ax.set_xlabel('Timecode')
    ax.set_ylabel('Mean Squared Error')
    ax.set_title(title)

    ax.grid(True, linestyle='--', alpha=0.7)
    ax.legend()
    plt.tight_layout()
    plt.close()
    return fig, anomaly_frames

def plot_mse_histogram(mse_values, title, anomaly_threshold, color='blue'):
    plt.figure(figsize=(16, 3), dpi=300)
    fig, ax = plt.subplots(figsize=(16, 3))

    ax.hist(mse_values, bins=100, edgecolor='black', color=color, alpha=0.7)
    ax.set_xlabel('Mean Squared Error')
    ax.set_ylabel('Number of Frames')
    ax.set_title(title)

    mean = np.mean(mse_values)
    std = np.std(mse_values)
    threshold = mean + anomaly_threshold * std

    ax.axvline(x=threshold, color='red', linestyle='--', linewidth=2)

    plt.tight_layout()
    plt.close()
    return fig

def plot_mse_heatmap(mse_values, title, df):
    plt.figure(figsize=(20, 3), dpi=300)
    fig, ax = plt.subplots(figsize=(20, 3))

    # Reshape MSE values to 2D array for heatmap
    mse_2d = mse_values.reshape(1, -1)

    # Create heatmap
    sns.heatmap(mse_2d, cmap='YlOrRd', cbar=False, ax=ax)

    # Set x-axis ticks to timecodes
    num_ticks = min(60, len(mse_values))
    tick_locations = np.linspace(0, len(mse_values) - 1, num_ticks).astype(int)
    
    # Ensure tick_locations are within bounds
    tick_locations = tick_locations[tick_locations < len(df)]
    
    tick_labels = [df['Timecode'].iloc[i] if i < len(df) else '' for i in tick_locations]

    ax.set_xticks(tick_locations)
    ax.set_xticklabels(tick_labels, rotation=90, ha='center', va='top')

    ax.set_title(title)

    # Remove y-axis labels
    ax.set_yticks([])

    plt.tight_layout()
    plt.close()
    return fig

def plot_posture(df, posture_scores, color='blue', anomaly_threshold=3):
    plt.figure(figsize=(16, 8), dpi=300)
    fig, ax = plt.subplots(figsize=(16, 8))

    df['Seconds'] = df['Timecode'].apply(
        lambda x: sum(float(t) * 60 ** i for i, t in enumerate(reversed(x.split(':')))))

    posture_data = [(frame, score) for frame, score in posture_scores.items() if score is not None]
    posture_frames, posture_scores = zip(*posture_data)

    # Create a new dataframe for posture data
    posture_df = pd.DataFrame({'Frame': posture_frames, 'Score': posture_scores})

    posture_df = posture_df.merge(df[['Frame', 'Seconds']], on='Frame', how='inner')

    ax.scatter(posture_df['Seconds'], posture_df['Score'], color=color, alpha=0.3, s=5)
    mean = posture_df['Score'].rolling(window=10).mean()
    ax.plot(posture_df['Seconds'], mean, color=color, linewidth=0.5)

    ax.set_xlabel('Timecode')
    ax.set_ylabel('Posture Score')
    ax.set_title("Body Posture Over Time")

    ax.grid(True, linestyle='--', alpha=0.7)

    max_seconds = df['Seconds'].max()
    num_ticks = 80
    tick_locations = np.linspace(0, max_seconds, num_ticks)
    tick_labels = [seconds_to_timecode(int(s)) for s in tick_locations]

    ax.set_xticks(tick_locations)
    ax.set_xticklabels(tick_labels, rotation=90, ha='center', fontsize=6)

    plt.tight_layout()
    plt.close()
    return fig

def filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames):
    # Ensure most_frequent_person_frames is a list
    if not isinstance(most_frequent_person_frames, (list, np.ndarray)):
        most_frequent_person_frames = [most_frequent_person_frames]

    # Ensure df and mse arrays have the same length
    min_length = min(len(df), len(mse_embeddings), len(mse_posture), len(mse_voice))
    df = df.iloc[:min_length].copy()
    mse_embeddings = mse_embeddings[:min_length]
    mse_posture = mse_posture[:min_length]
    mse_voice = mse_voice[:min_length]

    # Create a mask for the most frequent person frames
    mask = df['Frame'].isin(most_frequent_person_frames)

    # Pad mask to match the length of the video frames
    padded_mask = np.zeros(len(mse_embeddings), dtype=bool)
    padded_mask[:len(mask)] = mask

    # Apply the mask to filter the MSE arrays
    mse_embeddings_filtered = np.where(padded_mask, mse_embeddings, 0)
    mse_posture_filtered = np.where(padded_mask, mse_posture, 0)
    mse_voice_filtered = np.where(padded_mask, mse_voice, 0)

    return mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered

def normalize_mse(mse):
    return mse / np.max(mse) if np.max(mse) > 0 else mse

def pad_or_trim_array(arr, target_length):
    if len(arr) > target_length:
        return arr[:target_length]
    elif len(arr) < target_length:
        return np.pad(arr, (0, target_length - len(arr)), 'constant')
    return arr

def create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, fps, total_frames, width):
    frame_index = min(int(t * fps), len(mse_embeddings_filtered) - 1)
    
    # Normalize the MSE values
    mse_embeddings_norm = normalize_mse(mse_embeddings_filtered)
    mse_posture_norm = normalize_mse(mse_posture_filtered)
    mse_voice_norm = normalize_mse(mse_voice_filtered)

    # Ensure all arrays have the correct length
    mse_embeddings_norm = pad_or_trim_array(mse_embeddings_norm, total_frames)
    mse_posture_norm = pad_or_trim_array(mse_posture_norm, total_frames)
    mse_voice_norm = pad_or_trim_array(mse_voice_norm, total_frames)

    # Create a 3D array for the heatmap (height, width, channels)
    heatmap_height = 3  # Assuming you want 3 rows in your heatmap
    heatmap_frame = np.zeros((heatmap_height, width, 3), dtype=np.uint8)

    # Fill the heatmap frame with color based on MSE values
    heatmap_frame[0, :, 0] = (mse_embeddings_norm[frame_index] * 255).astype(np.uint8)  # Red channel for facial features
    heatmap_frame[1, :, 1] = (mse_posture_norm[frame_index] * 255).astype(np.uint8)    # Green channel for body posture
    heatmap_frame[2, :, 2] = (mse_voice_norm[frame_index] * 255).astype(np.uint8)      # Blue channel for voice

    return heatmap_frame
    

def create_video_with_heatmap(video_path, df, mse_embeddings, mse_posture, mse_voice, output_folder, desired_fps, most_frequent_person_frames):
    print(f"Creating heatmap video. Output folder: {output_folder}")
    
    os.makedirs(output_folder, exist_ok=True)
    
    output_filename = os.path.basename(video_path).rsplit('.', 1)[0] + '_heatmap.mp4'
    heatmap_video_path = os.path.join(output_folder, output_filename)
    
    print(f"Heatmap video will be saved at: {heatmap_video_path}")

    # Load the original video
    video = VideoFileClip(video_path)
    
    # Get video properties
    width, height = video.w, video.h
    total_frames = int(video.duration * video.fps)
    
    # Ensure all MSE arrays have the same length as total_frames
    mse_embeddings = np.interp(np.linspace(0, len(mse_embeddings) - 1, total_frames), 
                               np.arange(len(mse_embeddings)), mse_embeddings)
    mse_posture = np.interp(np.linspace(0, len(mse_posture) - 1, total_frames), 
                            np.arange(len(mse_posture)), mse_posture)
    mse_voice = np.interp(np.linspace(0, len(mse_voice) - 1, total_frames), 
                          np.arange(len(mse_voice)), mse_voice)

    print(f"Total frames: {total_frames}")
    print(f"mse_embeddings length: {len(mse_embeddings)}")
    print(f"mse_posture length: {len(mse_posture)}")
    print(f"mse_voice length: {len(mse_voice)}")
    
    # Filter MSE arrays for the most frequent person frames
    mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered = filter_mse_for_most_frequent_person(df, mse_embeddings, mse_posture, mse_voice, most_frequent_person_frames)
    
    def combine_video_and_heatmap(t):
        video_frame = video.get_frame(t)
        heatmap_frame = create_heatmap(t, mse_embeddings_filtered, mse_posture_filtered, mse_voice_filtered, video.fps, total_frames, width)
        heatmap_frame_resized = cv2.resize(heatmap_frame, (width, heatmap_frame.shape[0]))
        
        # Ensure both frames have the same number of channels
        if video_frame.shape[2] != heatmap_frame_resized.shape[2]:
            heatmap_frame_resized = cv2.cvtColor(heatmap_frame_resized, cv2.COLOR_RGB2BGR)
        
        combined_frame = np.vstack((video_frame, heatmap_frame_resized))
        return combined_frame

    final_clip = VideoClip(combine_video_and_heatmap, duration=video.duration)
    final_clip = final_clip.set_audio(video.audio)

    # Write the final video using x264 codec
    final_clip.write_videofile(heatmap_video_path, codec='libx264', audio_codec='aac', fps=video.fps, preset='medium', ffmpeg_params=['-crf', '23'])
    
    # Close the video clips
    video.close()
    final_clip.close()
    
    if os.path.exists(heatmap_video_path):
        print(f"Heatmap video created at: {heatmap_video_path}")
        print(f"Heatmap video size: {os.path.getsize(heatmap_video_path)} bytes")
        return heatmap_video_path
    else:
        print(f"Failed to create heatmap video at: {heatmap_video_path}")
        return None

def plot_correlation_heatmap(mse_embeddings, mse_posture, mse_voice):
    data = np.vstack((mse_embeddings, mse_posture, mse_voice)).T
    df = pd.DataFrame(data, columns=["Facial Features", "Body Posture", "Voice"])
    corr = df.corr()

    plt.figure(figsize=(10, 8), dpi=300)

    heatmap = sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
    plt.title('Correlation Heatmap of MSEs')
    plt.tight_layout()
    return plt.gcf()