Spaces:
Sleeping
Sleeping
File size: 5,081 Bytes
5079645 e3641b1 5079645 4daa026 e3641b1 5079645 4daa026 5079645 e3641b1 5079645 e3641b1 d34196f e3641b1 d34196f e3641b1 5079645 e3641b1 5079645 e3641b1 5079645 e3641b1 5079645 e3641b1 4daa026 e3641b1 4daa026 e3641b1 4daa026 e3641b1 5079645 e3641b1 5079645 e3641b1 5079645 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
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,
show_arrows,
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_teacher=detection_result_teacher,
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,
show_arrows=show_arrows,
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, arrow) for (log, arrow) 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
|