add arrows

app.py

@@ -59,6 +59,9 @@ with gr.Blocks() as demo:
 
             trigger_state = gr.Radio(value="one", choices=["one", "two", "three"], label="Trigger Count")
 
+            gr.Markdown("#### Plot arrows:")
+            show_arrows = gr.Checkbox(label="If True, arrows will be plotted on the video")
+
         input_teacher = gr.Video(show_share_button=False, show_download_button=False, sources=["upload"], label="Teacher's Video")
         input_student = gr.Video(show_share_button=False, show_download_button=False, sources=["upload"], label="Student's Video")
 
@@ -141,7 +144,7 @@ with gr.Blocks() as demo:
 
     gr_button.click(
         fn=video_identity,
-        inputs=[dtw_mean, dtw_filter, angles_sensitive, angles_common, angles_insensitive, trigger_state, input_teacher, input_student],
+        inputs=[dtw_mean, dtw_filter, angles_sensitive, angles_common, angles_insensitive, trigger_state, show_arrows, input_teacher, input_student],
         outputs=[output_merged, general_log, text_log]
     )
 

main_func.py

@@ -11,7 +11,7 @@ from utils import (
     get_dtw_mean_path,
     generate_output_video,
     generate_log,
-    write_log,
+    write_log
 )
 
 from config import (
@@ -30,6 +30,7 @@ def video_identity(
         angles_common, 
         angles_insensitive, 
         trigger_state,
+        show_arrows,
         video_teacher, 
         video_student
 ):
@@ -91,6 +92,7 @@ def video_identity(
     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,
@@ -103,12 +105,13 @@ def video_identity(
             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) for log in text_info_summary])
+        all_text_summaries.extend([(log, idx, arrow) for (log, arrow) in text_info_summary])
         
 # ======================================================================================
 # create files for downloading and displaying.

utils.py

@@ -10,6 +10,7 @@ from datetime import timedelta
 from scipy.signal import savgol_filter
 from scipy.stats import mstats
 
+
 def predict_keypoints_vitpose(
         video_path, 
         model_path, 
@@ -229,6 +230,7 @@ def get_dtw_mean_path(serieses_teacher, serieses_student, dtw_mean, dtw_filter):
 
 
 def modify_student_frame(
+    detection_result_teacher,
     detection_result_student,
     detection_result_teacher_angles,
     detection_result_student_angles,
@@ -241,9 +243,14 @@ def modify_student_frame(
     previously_trigered,
     previously_trigered_2,
     triger_state,
+    show_arrows,
     text_dictionary,
 ):
-    
+    arrows_bgr = (175, 75, 190)
+    arrows_sz = 3
+    skeleton_bgr = (0, 0, 255)
+    skeleton_sz = 3
+
     frame_copy = video_student[alignment_frames[1]]
     frame_teacher_copy = video_teacher[alignment_frames[0]]
     frame_errors = np.abs(detection_result_teacher_angles[alignment_frames[0]] - detection_result_student_angles[alignment_frames[1]])
@@ -272,22 +279,27 @@ def modify_student_frame(
 
             if check_threshold and check_certain and relevant_plane:
 
-                point1 = detection_result_student[:, :, :-1][alignment_frames[1]][connection[0]]
-                point2 = detection_result_student[:, :, :-1][alignment_frames[1]][connection[1]]
-
-                point1 = np.array(point1).astype(int)
-                point2 = np.array(point2).astype(int)
-
-                point1 = [point1[1],  point1[0]]
-                point2 = [ point2[1],  point2[0]]
+                point1, point2, point2_t = align_points(
+                    detection_result_student, 
+                    detection_result_teacher, 
+                    alignment_frames, 
+                    edge
+                )
 
+                arrow = get_arrow_direction(point2, point2_t)
 
                 if triger_state == "one":
 
-                    _ = cv2.line(frame_copy, point1, point2, (0, 0, 255), 10)
+                    _ = cv2.line(frame_copy, point1, point2, skeleton_bgr, skeleton_sz)
+
+                    if show_arrows:
+                        _ = cv2.arrowedLine(frame_copy, point2, point2_t, arrows_bgr, arrows_sz) 
 
                     if (connection[0], connection[1]) in text_dictionary:
-                        text_info.append(text_dictionary[(connection[0], connection[1])])
+                        text_info.append((text_dictionary[(connection[0], connection[1])], arrow))
+                                
+                    if (connection[1], connection[0]) in text_dictionary:
+                        text_info.append((text_dictionary[(connection[1], connection[0])], arrow))
 
                 if triger_state == "two":
 
@@ -295,10 +307,16 @@ def modify_student_frame(
 
                     if (connection[0], connection[1]) in previously_trigered:
 
-                        _ = cv2.line(frame_copy, point1, point2, (0, 0, 255), 10)
+                        _ = cv2.line(frame_copy, point1, point2, skeleton_bgr, skeleton_sz)
+
+                        if show_arrows:
+                            _ = cv2.arrowedLine(frame_copy, point2, point2_t, arrows_bgr, arrows_sz) 
 
                         if (connection[0], connection[1]) in text_dictionary:
-                            text_info.append(text_dictionary[(connection[0], connection[1])])
+                            text_info.append((text_dictionary[(connection[0], connection[1])], arrow))
+                                
+                        if (connection[1], connection[0]) in text_dictionary:
+                            text_info.append((text_dictionary[(connection[1], connection[0])], arrow))
 
                 if triger_state == "three":
 
@@ -310,10 +328,16 @@ def modify_student_frame(
 
                         if (connection[0], connection[1]) in previously_trigered_2:
 
-                            _ = cv2.line(frame_copy, point1, point2, (0, 0, 255), 10)
+                            _ = cv2.line(frame_copy, point1, point2, skeleton_bgr, skeleton_sz)
+
+                            if show_arrows:
+                                _ = cv2.arrowedLine(frame_copy, point2, point2_t, arrows_bgr, arrows_sz) 
 
                             if (connection[0], connection[1]) in text_dictionary:
-                                text_info.append(text_dictionary[(connection[0], connection[1])])
+                                text_info.append((text_dictionary[(connection[0], connection[1])], arrow))
+
+                            if (connection[1], connection[0]) in text_dictionary:
+                                text_info.append((text_dictionary[(connection[1], connection[0])], arrow))
            
     return frame_copy, frame_teacher_copy, list(set(trigered_connections)), list(set(trigered_connections2)), text_info
 
@@ -412,9 +436,9 @@ def generate_log(all_text_summaries):
 
     general_summary = []
     for log in all_text_summaries_clean:
-        comment, frame = log
+        comment, frame, arrow = log
         total_seconds = frame / 30
-        general_summary.append(f"{comment}. Video time: {str(timedelta(seconds=total_seconds))[3:-4]}")
+        general_summary.append(f"{comment}. Direction: {arrow}. Video time: {str(timedelta(seconds=total_seconds))[3:-4]}")
 
     general_summary = "\n".join(general_summary)
 
@@ -465,3 +489,69 @@ Error logs:
     print(f"log {log_path} was created.")
 
     return log_path
+
+
+def angle_between(v1, v2):
+    return np.arctan2(v2[1], v2[0]) - np.arctan2(v1[1], v1[0])
+
+
+def align_points(detection_result_student, detection_result_teacher, alignment_frames, edge):
+
+    point0 = detection_result_student[alignment_frames[1], edge[0], :-1].astype(int)[::-1]
+    point1 = detection_result_student[alignment_frames[1], edge[1], :-1].astype(int)[::-1]
+    point2 = detection_result_student[alignment_frames[1], edge[2], :-1].astype(int)[::-1]
+
+    point0_t = detection_result_teacher[alignment_frames[0], edge[0], :-1].astype(int)[::-1]
+    point1_t = detection_result_teacher[alignment_frames[0], edge[1], :-1].astype(int)[::-1]
+    point2_t = detection_result_teacher[alignment_frames[0], edge[2], :-1].astype(int)[::-1]
+
+    translation = point0 - point0_t
+
+    point0_t += translation
+    point1_t += translation
+    point2_t += translation
+
+    BsA = point1 - point0
+    BtA = point1_t - point0
+
+    theta = angle_between(BtA, BsA)
+
+    R = np.array([
+        [np.cos(theta), -np.sin(theta)],
+        [np.sin(theta), np.cos(theta)]
+    ])
+
+    point1_t = np.dot(R, (point1_t - point0).T).T + point0
+    point2_t = np.dot(R, (point2_t - point0).T).T + point0
+
+    point2_t = point2_t.astype(int)
+
+    return point1, point2, point2_t
+
+
+def get_arrow_direction(A, B):
+
+    translation_vector = B - A
+    angle_deg = np.degrees(np.arctan2(translation_vector[0], translation_vector[1]))
+
+    match angle_deg:
+        case angle if -22.5 <= angle < 22.5:
+            arrow = "⬆"
+        case angle if 22.5 <= angle < 67.5:
+            arrow = "⬈"
+        case angle if 67.5 <= angle < 112.5:
+            arrow = "➡"
+        case angle if 112.5 <= angle < 157.5:
+            arrow = "⬊"
+        case angle if 157.5 <= angle or angle < -157.5:
+            arrow = "⬇"
+        case angle if -157.5 <= angle < -112.5:
+            arrow = "⬋"
+        case angle if -112.5 <= angle < -67.5:
+            arrow = "⬅"
+        case angle if -67.5 <= angle < -22.5:
+            arrow = "⬉"
+        case _:
+            arrow = ""
+
+    return arrow