main_func.py

import datetime

from utils import (
    predict_keypoints_vitpose, 
    get_edge_groups, 
    get_series,
    z_score_normalization,
    modify_student_frame,
    get_video_frames,
    check_and_download_models,
    get_dtw_mean_path,
    generate_output_video,
    generate_log,
    write_log,
)

from config import (
    CONNECTIONS_VIT_FULL,
    CONNECTIONS_FOR_ERROR,
    EDGE_GROUPS_FOR_ERRORS,
    EDGE_GROUPS_FOR_SUMMARY,
    get_thresholds
)


def video_identity(
        dtw_mean, 
        dtw_filter,
        angles_sensitive, 
        angles_common, 
        angles_insensitive, 
        trigger_state,
        video_teacher, 
        video_student
):

# ======================================================================================
# This part is responsible for keypoints detection on two videos.
# ======================================================================================
    check_and_download_models()

    detection_result_teacher = predict_keypoints_vitpose(
    video_path=video_teacher,
    model_path="models/vitpose-b-wholebody.pth",
    model_name="b",
    detector_path="models/yolov8s.pt"
    )

    detection_result_student = predict_keypoints_vitpose(
        video_path=video_student,
        model_path="models/vitpose-b-wholebody.pth",
        model_name="b",
        detector_path="models/yolov8s.pt"
    )


# ======================================================================================
# Here we perform transformations of keypoints to angles between keypoints and normalize them.
# ======================================================================================
    detection_result_teacher_angles = get_series(detection_result_teacher[:, :,:-1], EDGE_GROUPS_FOR_ERRORS).T
    detection_result_student_angles = get_series(detection_result_student[:, :,:-1], EDGE_GROUPS_FOR_ERRORS).T


    edge_groups_for_dtw = get_edge_groups(CONNECTIONS_VIT_FULL)
    serieses_teacher = get_series(detection_result_teacher[:, :,:-1], edge_groups_for_dtw)
    serieses_student = get_series(detection_result_student[:, :,:-1], edge_groups_for_dtw)

    serieses_teacher = z_score_normalization(serieses_teacher)
    serieses_student = z_score_normalization(serieses_student)


# ======================================================================================
# Finding of frame to frame mean alignment with DTW algorithm.
# ======================================================================================
    alignments = get_dtw_mean_path(serieses_teacher, serieses_student, dtw_mean, dtw_filter)


# ======================================================================================
# Adding visual marks on student's video, speed alignment and error log generation.
# ======================================================================================
    trigger_1 = []
    trigger_2 = []
    save_teacher_frames = []
    save_student_frames = []
    all_text_summaries = []

    video_teacher_loaded = get_video_frames(video_teacher)
    video_student_loaded = get_video_frames(video_student)
    threshouds_for_errors = get_thresholds(angles_sensitive, angles_common, angles_insensitive)

    for idx, alignment in enumerate(alignments):

        frame_student_out, frame_teacher_out, trigger_1, trigger_2, text_info_summary = modify_student_frame(
            detection_result_student=detection_result_student,
            detection_result_teacher_angles=detection_result_teacher_angles,
            detection_result_student_angles=detection_result_student_angles,
            video_teacher=video_teacher_loaded,
            video_student=video_student_loaded,
            alignment_frames=alignment,
            edge_groups=EDGE_GROUPS_FOR_ERRORS,
            connections=CONNECTIONS_FOR_ERROR,
            thresholds=threshouds_for_errors,
            previously_trigered=trigger_1,
            previously_trigered_2=trigger_2,
            triger_state=trigger_state,
            text_dictionary=EDGE_GROUPS_FOR_SUMMARY
        )

        save_teacher_frames.append(frame_teacher_out)
        save_student_frames.append(frame_student_out)
        all_text_summaries.extend([(log, idx) for log in text_info_summary])
        
# ======================================================================================
# create files for downloading and displaying.
# ======================================================================================
   
    timestamp_str = datetime.datetime.now().strftime("%Y_%m-%d_%H_%M_%S")

    video_path = generate_output_video(save_teacher_frames, save_student_frames, timestamp_str)
    
    general_summary = generate_log(all_text_summaries)

    log_path = write_log(
        timestamp_str, 
        dtw_mean, 
        dtw_filter, 
        angles_sensitive, 
        angles_common, 
        angles_insensitive,
        trigger_state,
        general_summary
    )

    return video_path, general_summary, log_path