video_processing.py

import os
import cv2
import numpy as np
from moviepy.editor import VideoFileClip
import tempfile
import time
from PIL import Image, ImageDraw, ImageFont
import math
from face_analysis import get_face_embedding, cluster_faces, organize_faces_by_person
from pose_analysis import pose, calculate_posture_score, draw_pose_landmarks
from voice_analysis import get_speaker_embeddings, align_voice_embeddings, extract_audio_from_video, diarize_speakers
from anomaly_detection import anomaly_detection
from visualization import plot_mse, plot_mse_histogram, plot_mse_heatmap, plot_stacked_mse_heatmaps
from utils import frame_to_timecode
import pandas as pd
from facenet_pytorch import MTCNN
import torch
import mediapipe as mp

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.98, 0.98, 0.98], min_face_size=200, post_process=False)

def extract_frames(video_path, output_folder, desired_fps, progress_callback=None):
    os.makedirs(output_folder, exist_ok=True)
    clip = VideoFileClip(video_path)
    original_fps = clip.fps
    duration = clip.duration
    total_frames = int(duration * original_fps)
    step = max(1, original_fps / desired_fps)
    total_frames_to_extract = int(total_frames / step)

    frame_count = 0
    for t in np.arange(0, duration, step / original_fps):
        frame = clip.get_frame(t)
        cv2.imwrite(os.path.join(output_folder, f"frame_{frame_count:04d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
        frame_count += 1
        if progress_callback:
            progress = min(100, (frame_count / total_frames_to_extract) * 100)
            progress_callback(progress, f"Extracting frame")
        if frame_count >= total_frames_to_extract:
            break
    clip.close()
    return frame_count, original_fps


def is_frontal_face(face, landmarks):
    if landmarks is None:
        return False
    
    left_eye = landmarks[0]
    right_eye = landmarks[1]
    nose = landmarks[2]
    
    eye_angle = np.degrees(np.arctan2(right_eye[1] - left_eye[1], right_eye[0] - left_eye[0]))
    
    eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
    nose_deviation = abs(nose[0] - eye_center[0]) / face.shape[1]
    
    return abs(eye_angle) < 10 and nose_deviation < 0.1

def process_frames(frames_folder, faces_folder, frame_count, progress):
    embeddings_by_frame = {}
    posture_scores_by_frame = {}
    posture_landmarks_by_frame = {}
    face_paths = []
    frame_files = sorted([f for f in os.listdir(frames_folder) if f.endswith('.jpg')])

    for i, frame_file in enumerate(frame_files):
        frame_num = int(frame_file.split('_')[1].split('.')[0])
        frame_path = os.path.join(frames_folder, frame_file)
        frame = cv2.cvtColor(cv2.imread(frame_path), cv2.COLOR_BGR2RGB)

        if frame is not None:
            posture_score, posture_landmarks = calculate_posture_score(frame)
            posture_scores_by_frame[frame_num] = posture_score
            posture_landmarks_by_frame[frame_num] = posture_landmarks

            boxes, probs, landmarks = mtcnn.detect(frame, landmarks=True)

            if boxes is not None and len(boxes) > 0 and probs[0] >= 0.99:
                x1, y1, x2, y2 = [int(b) for b in boxes[0]]
                face = frame[y1:y2, x1:x2]

                if face.size > 0 and is_frontal_face(face, landmarks[0]):
                    face_resized = cv2.resize(face, (160, 160))
                    output_path = os.path.join(faces_folder, f"frame_{frame_num}_face.jpg")
                    cv2.imwrite(output_path, cv2.cvtColor(face_resized, cv2.COLOR_RGB2BGR))
                    face_paths.append(output_path)
                    embedding = get_face_embedding(face_resized)
                    embeddings_by_frame[frame_num] = embedding

        progress((i + 1) / len(frame_files), f"Processing frame {i + 1} of {len(frame_files)}")

    return embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, face_paths


def process_video(video_path, anomaly_threshold, desired_fps, progress=None):
    start_time = time.time()
    output_folder = "output"
    os.makedirs(output_folder, exist_ok=True)

    with tempfile.TemporaryDirectory() as temp_dir:
        faces_folder = os.path.join(temp_dir, 'faces')
        organized_faces_folder = os.path.join(temp_dir, 'organized_faces')
        os.makedirs(faces_folder, exist_ok=True)
        os.makedirs(organized_faces_folder, exist_ok=True)

        clip = VideoFileClip(video_path)
        video_duration = clip.duration
        clip.close()

        progress(0, "Starting frame extraction")
        frames_folder = os.path.join(temp_dir, 'extracted_frames')

        def extraction_progress(percent, message):
            progress(percent / 100, f"Extracting frames")

        frame_count, original_fps = extract_frames(video_path, frames_folder, desired_fps, extraction_progress)

        progress(1, "Frame extraction complete")
        progress(0.3, "Processing frames")
        embeddings_by_frame, posture_scores_by_frame, posture_landmarks_by_frame, face_paths = process_frames(
            frames_folder, faces_folder,
            frame_count,
            progress)

        if not face_paths:
            raise ValueError("No faces were extracted from the video.")

        progress(0.6, "Clustering faces")
        embeddings = [embedding for _, embedding in embeddings_by_frame.items()]
        clusters = cluster_faces(embeddings)
        num_clusters = len(set(clusters))

        # Adding the 'Cluster' column to the DataFrame
        cluster_by_frame = {frame_num: cluster for frame_num, cluster in zip(embeddings_by_frame.keys(), clusters)}

        progress(0.65, "Organizing faces")
        organize_faces_by_person(embeddings_by_frame, clusters, faces_folder, organized_faces_folder)

        progress(0.7, "Saving person data")
        df, largest_cluster = save_person_data(embeddings_by_frame, clusters, desired_fps,
                                                      original_fps, temp_dir, video_duration)

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

        progress(0.75, "Getting face samples")
        face_samples = get_all_face_samples(organized_faces_folder, output_folder, largest_cluster)

        progress(0.8, "Performing voice analysis")
        audio_path = extract_audio_from_video(video_path)
        diarization, most_frequent_speaker = diarize_speakers(audio_path)
        voice_embeddings, audio_duration = get_speaker_embeddings(audio_path, diarization, most_frequent_speaker)
        aligned_voice_embeddings = align_voice_embeddings(voice_embeddings, frame_count, original_fps, audio_duration)

        progress(0.85, "Performing anomaly detection")
        embedding_columns = [col for col in df.columns if col.startswith('Raw_Embedding_')]

        X_embeddings = df[embedding_columns].values
        X_posture = np.array([posture_scores_by_frame.get(frame, None) for frame in df['Frame']])
        X_posture = X_posture[X_posture != None].reshape(-1, 1)
        X_voice = np.array(aligned_voice_embeddings)

        if len(X_voice) > len(X_embeddings):
            X_voice = X_voice[:len(X_embeddings)]
        elif len(X_voice) < len(X_embeddings):
            padding = np.zeros((len(X_embeddings) - len(X_voice), X_voice.shape[1]))
            X_voice = np.vstack((X_voice, padding))

        try:
            if len(X_posture) == 0:
                raise ValueError("No valid posture data found")

            mse_embeddings, mse_posture, mse_voice = anomaly_detection(X_embeddings, X_posture, X_voice)

            progress(0.9, "Generating graphs")
            mse_plot_embeddings, anomaly_frames_embeddings = plot_mse(df, mse_embeddings, "Facial Features",
                                                                      color='navy',
                                                                      anomaly_threshold=anomaly_threshold)

            mse_histogram_embeddings = plot_mse_histogram(mse_embeddings, "MSE Distribution: Facial Features",
                                                          anomaly_threshold, color='navy')

            mse_plot_posture, anomaly_frames_posture = plot_mse(df, mse_posture, "Body Posture",
                                                                color='purple',
                                                                anomaly_threshold=anomaly_threshold)

            mse_histogram_posture = plot_mse_histogram(mse_posture, "MSE Distribution: Body Posture",
                                                       anomaly_threshold, color='purple')

            mse_plot_voice, anomaly_frames_voice = plot_mse(df, mse_voice, "Voice",
                                                            color='green',
                                                            anomaly_threshold=anomaly_threshold)

            mse_histogram_voice = plot_mse_histogram(mse_voice, "MSE Distribution: Voice",
                                                     anomaly_threshold, color='green')

            mse_heatmap_embeddings = plot_mse_heatmap(mse_embeddings, "Facial Features MSE Heatmap", df)
            mse_heatmap_posture = plot_mse_heatmap(mse_posture, "Body Posture MSE Heatmap", df)
            mse_heatmap_voice = plot_mse_heatmap(mse_voice, "Voice MSE Heatmap", df)

            stacked_heatmap = plot_stacked_mse_heatmaps(mse_embeddings, mse_posture, mse_voice, df, "Combined MSE Heatmaps")

            progress(0.95, "Finishing generating graphs")

        except Exception as e:
            print(f"Error details: {str(e)}")
            import traceback
            traceback.print_exc()
            return (f"Error in video processing: {str(e)}",) + (None,) * 26

        progress(1.0, "Preparing results")
        results = f"Number of persons detected: {num_clusters}\n\n"
        results += "Breakdown:\n"
        for cluster_id in range(num_clusters):
            face_count = len([c for c in clusters if c == cluster_id])
            results += f"Person {cluster_id + 1}: {face_count} face frames\n"

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

        def add_timecode_to_image(image, timecode):
            img_pil = Image.fromarray(image)
            draw = ImageDraw.Draw(img_pil)
            font = ImageFont.load_default()
            draw.text((10, 10), timecode, (255, 0, 0), font=font)
            return np.array(img_pil)

        anomaly_faces_embeddings = []
        for frame in anomaly_frames_embeddings:
            face_path = os.path.join(faces_folder, f"frame_{frame}_face.jpg")
            if os.path.exists(face_path):
                face_img = cv2.imread(face_path)
                if face_img is not None:
                    face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
                    timecode = df[df['Frame'] == frame]['Timecode'].iloc[0]
                    face_img_with_timecode = add_timecode_to_image(face_img, timecode)
                    anomaly_faces_embeddings.append(face_img_with_timecode)

        anomaly_frames_posture_images = []
        for frame in anomaly_frames_posture:
            frame_path = os.path.join(frames_folder, f"frame_{frame:04d}.jpg")
            if os.path.exists(frame_path):
                frame_img = cv2.imread(frame_path)
                if frame_img is not None:
                    frame_img = cv2.cvtColor(frame_img, cv2.COLOR_BGR2RGB)
                    pose_results = pose.process(frame_img)
                    if pose_results.pose_landmarks:
                        frame_img = draw_pose_landmarks(frame_img, pose_results.pose_landmarks)
                    timecode = df[df['Frame'] == frame]['Timecode'].iloc[0]
                    frame_img_with_timecode = add_timecode_to_image(frame_img, timecode)
                    anomaly_frames_posture_images.append(frame_img_with_timecode)

        return (
            execution_time,
            results,
            df,
            mse_embeddings,
            mse_posture,
            mse_voice,
            mse_plot_embeddings,
            mse_plot_posture,
            mse_plot_voice,
            mse_histogram_embeddings,
            mse_histogram_posture,
            mse_histogram_voice,
            mse_heatmap_embeddings,
            mse_heatmap_posture,
            mse_heatmap_voice,
            face_samples["most_frequent"],
            anomaly_faces_embeddings,
            anomaly_frames_posture_images,
            faces_folder,
            frames_folder,
            stacked_heatmap
              
        )


def save_person_data(embeddings_by_frame, clusters, desired_fps, original_fps, output_folder, video_duration):
    person_data = {}

    for (frame_num, embedding), cluster in zip(embeddings_by_frame.items(), clusters):
        if cluster not in person_data:
            person_data[cluster] = []
        person_data[cluster].append((frame_num, embedding))

    largest_cluster = max(person_data, key=lambda k: len(person_data[k]))

    data = person_data[largest_cluster]
    data.sort(key=lambda x: x[0])
    frames, embeddings = zip(*data)

    embeddings_array = np.array(embeddings)
    np.save(os.path.join(output_folder, 'face_embeddings.npy'), embeddings_array)

    total_frames = max(frames)
    timecodes = [frame_to_timecode(frame, total_frames, video_duration) for frame in frames]

    df_data = {
        'Frame': frames,
        'Timecode': timecodes,
        'Embedding_Index': range(len(embeddings))
    }

    for i in range(len(embeddings[0])):
        df_data[f'Raw_Embedding_{i}'] = [embedding[i] for embedding in embeddings]

    df = pd.DataFrame(df_data)

    return df, largest_cluster

def get_all_face_samples(organized_faces_folder, output_folder, largest_cluster, max_samples=200):
    face_samples = {"most_frequent": [], "others": []}
    for cluster_folder in sorted(os.listdir(organized_faces_folder)):
        if cluster_folder.startswith("person_"):
            person_folder = os.path.join(organized_faces_folder, cluster_folder)
            face_files = sorted([f for f in os.listdir(person_folder) if f.endswith('.jpg')])
            if face_files:
                cluster_id = int(cluster_folder.split('_')[1])
                if cluster_id == largest_cluster:
                    for i, sample in enumerate(face_files[:max_samples]):
                        face_path = os.path.join(person_folder, sample)
                        output_path = os.path.join(output_folder, f"face_sample_most_frequent_{i:04d}.jpg")
                        face_img = cv2.cvtColor(cv2.imread(face_path), cv2.COLOR_BGR2RGB)
                        if face_img is not None:
                            small_face = cv2.resize(face_img, (160, 160))
                            cv2.imwrite(output_path, cv2.cvtColor(small_face, cv2.COLOR_RGB2BGR))
                            face_samples["most_frequent"].append(output_path)
                        if len(face_samples["most_frequent"]) >= max_samples:
                            break
                else:
                    remaining_samples = max_samples - len(face_samples["others"])
                    if remaining_samples > 0:
                        for i, sample in enumerate(face_files[:remaining_samples]):
                            face_path = os.path.join(person_folder, sample)
                            output_path = os.path.join(output_folder, f"face_sample_other_{cluster_id:02d}_{i:04d}.jpg")
                            face_img = cv2.cvtColor(cv2.imread(face_path), cv2.COLOR_BGR2RGB)
                            if face_img is not None:
                                small_face = cv2.resize(face_img, (160, 160))
                                cv2.imwrite(output_path, cv2.cvtColor(small_face, cv2.COLOR_RGB2BGR))
                                face_samples["others"].append(output_path)
                            if len(face_samples["others"]) >= max_samples:
                                break
    return face_samples