Update main.py

main.py

@@ -1,8 +1,10 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
 import copy
 import math
 import argparse
 
-import cv2 as cv
+import cv2
 import numpy as np
 import mediapipe as mp
 
@@ -51,9 +53,10 @@ def main():
 
     rev_color = args.rev_color
 
-    cap = cv.VideoCapture(cap_device)
-    cap.set(cv.CAP_PROP_FRAME_WIDTH, cap_width)
-    cap.set(cv.CAP_PROP_FRAME_HEIGHT, cap_height)
+    cap = cv2.VideoCapture(cap_device)
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, cap_width)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cap_height)
+
 
     mp_pose = mp.solutions.pose
     pose = mp_pose.Pose(
@@ -64,14 +67,15 @@ def main():
     )
 
     cvFpsCalc = CvFpsCalc(buffer_len=10)
+    color = (100, 33, 3)
+    bg_color = (255, 255, 255)
 
-
-    if rev_color:
-        color = (255, 255, 255)
-        bg_color = (100, 33, 3)
-    else:
-        color = (100, 33, 3)
-        bg_color = (255, 255, 255)
+    # if rev_color:
+    #     color = (255, 255, 255)
+    #     bg_color = (100, 33, 3)
+    # else:
+    #     color = (100, 33, 3)
+    #     bg_color = (255, 255, 255)
 
     while True:
         display_fps = cvFpsCalc.get()
@@ -79,19 +83,17 @@ def main():
         ret, image = cap.read()
         if not ret:
             break
-        image = cv.flip(image, 1)  
+        image = cv2.flip(image, 1)
         debug_image01 = copy.deepcopy(image)
         debug_image02 = np.zeros((image.shape[0], image.shape[1], 3), np.uint8)
-        cv.rectangle(debug_image02, (0, 0), (image.shape[1], image.shape[0]),
+        cv2.rectangle(debug_image02, (0, 0), (image.shape[1], image.shape[0]),
                     bg_color,
                     thickness=-1)
 
-        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         results = pose.process(image)
 
-
         if results.pose_landmarks is not None:
-
             debug_image01 = draw_landmarks(
                 debug_image01,
                 results.pose_landmarks,
@@ -103,22 +105,20 @@ def main():
                 bg_color=bg_color,
             )
 
-        cv.putText(debug_image01, "FPS:" + str(display_fps), (10, 30),
-                   cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2, cv.LINE_AA)
-        cv.putText(debug_image02, "FPS:" + str(display_fps), (10, 30),
-                   cv.FONT_HERSHEY_SIMPLEX, 1.0, color, 2, cv.LINE_AA)
+        cv2.putText(debug_image01, "FPS:" + str(display_fps), (10, 30),
+                   cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2, cv2.LINE_AA)
+        cv2.putText(debug_image02, "FPS:" + str(display_fps), (10, 30),
+                   cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2, cv2.LINE_AA)
 
-
-        key = cv.waitKey(1)
-        if key == 27:  
+        key = cv2.waitKey(1)
+        if key == 27:
             break
 
-
-        cv.imshow('Tokyo2020 Debug', debug_image01)
-        cv.imshow('Tokyo2020 Pictogram', debug_image02)
+        cv2.imshow('Debug', debug_image01)
+        cv2.imshow('Pictogram', debug_image02)
 
     cap.release()
-    cv.destroyAllWindows()
+    cv2.destroyAllWindows()
 
 
 def draw_stick_figure(
@@ -130,7 +130,6 @@ def draw_stick_figure(
 ):
     image_width, image_height = image.shape[1], image.shape[0]
 
-
     landmark_point = []
     for index, landmark in enumerate(landmarks.landmark):
         landmark_x = min(int(landmark.x * image_width), image_width - 1)
@@ -139,7 +138,6 @@ def draw_stick_figure(
         landmark_point.append(
             [index, landmark.visibility, (landmark_x, landmark_y), landmark_z])
 
-
     right_leg = landmark_point[23]
     left_leg = landmark_point[24]
     leg_x = int((right_leg[2][0] + left_leg[2][0]) / 2)
@@ -148,12 +146,10 @@ def draw_stick_figure(
     landmark_point[23][2] = (leg_x, leg_y)
     landmark_point[24][2] = (leg_x, leg_y)
 
-
     sorted_landmark_point = sorted(landmark_point,
                                    reverse=True,
                                    key=lambda x: x[3])
 
-
     (face_x, face_y), face_radius = min_enclosing_face_circle(landmark_point)
 
     face_x = int(face_x)
@@ -164,21 +160,20 @@ def draw_stick_figure(
     stick_radius02 = int(stick_radius01 * (3 / 4))
     stick_radius03 = int(stick_radius02 * (3 / 4))
 
-
     draw_list = [
-        11, 
-        12,  
-        23,  
-        24,  
+        11,
+        12,
+        23,
+        24,
     ]
 
 
-    cv.rectangle(image, (0, 0), (image_width, image_height),
+    cv2.rectangle(image, (0, 0), (image_width, image_height),
                  bg_color,
                  thickness=-1)
 
 
-    cv.circle(image, (face_x, face_y), face_radius, color, -1)
+    cv2.circle(image, (face_x, face_y), face_radius, color, -1)
 
 
     for landmark_info in sorted_landmark_point:
@@ -224,7 +219,7 @@ def min_enclosing_face_circle(landmark_point):
         ]
         landmark_array = np.append(landmark_array, np_landmark_point, axis=0)
 
-    center, radius = cv.minEnclosingCircle(points=landmark_array)
+    center, radius = cv2.minEnclosingCircle(points=landmark_array)
 
     return center, radius
 
@@ -238,8 +233,8 @@ def draw_stick(
         color=(100, 33, 3),
         bg_color=(255, 255, 255),
 ):
-    cv.circle(image, point01, point01_radius, color, -1)
-    cv.circle(image, point02, point02_radius, color, -1)
+    cv2.circle(image, point01, point01_radius, color, -1)
+    cv2.circle(image, point02, point02_radius, color, -1)
 
     draw_list = []
     for index in range(2):
@@ -257,7 +252,7 @@ def draw_stick(
         draw_list.append([point_x, point_y])
 
     points = np.array((draw_list[0], draw_list[1], draw_list[3], draw_list[2]))
-    cv.fillConvexPoly(image, points=points, color=color)
+    cv2.fillConvexPoly(image, points=points, color=color)
 
     return image
 
@@ -265,7 +260,7 @@ def draw_stick(
 def draw_landmarks(
     image,
     landmarks,
-    # upper_body_only,
+
     visibility_th=0.5,
 ):
     image_width, image_height = image.shape[1], image.shape[0]
@@ -281,207 +276,211 @@ def draw_landmarks(
         if landmark.visibility < visibility_th:
             continue
 
-        if index == 0:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 1: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 2: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 3:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 4:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 5:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 6:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 7:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 8:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 9: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 10: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 11: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 12: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 13: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 14: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 15:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 16: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 17: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 18:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 19:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 20: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 21: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 22:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 23: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 24:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 25:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 26:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 27: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 28: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 29: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 30:  
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 31: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
-        if index == 32: 
-            cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 0:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 1:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 2:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 3:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 4:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 5:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 6:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 7:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 8:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 9:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 10:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 11:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 12:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 13:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 14:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 15:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 16:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 17:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 18:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 19:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 20:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 21:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 22:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 23:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 24:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 25:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 26:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 27:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 28:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 29:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 30:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 31:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
+        if index == 32:
+            cv2.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2)
 
         # if not upper_body_only:
         if True:
-            cv.putText(image, "z:" + str(round(landmark_z, 3)),
+            cv2.putText(image, "z:" + str(round(landmark_z, 3)),
                        (landmark_x - 10, landmark_y - 10),
-                       cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1,
-                       cv.LINE_AA)
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1,
+                       cv2.LINE_AA)
+
 
     if landmark_point[1][0] > visibility_th and landmark_point[2][
             0] > visibility_th:
-        cv.line(image, landmark_point[1][1], landmark_point[2][1],
+        cv2.line(image, landmark_point[1][1], landmark_point[2][1],
                 (0, 255, 0), 2)
     if landmark_point[2][0] > visibility_th and landmark_point[3][
             0] > visibility_th:
-        cv.line(image, landmark_point[2][1], landmark_point[3][1],
+        cv2.line(image, landmark_point[2][1], landmark_point[3][1],
                 (0, 255, 0), 2)
 
+
     if landmark_point[4][0] > visibility_th and landmark_point[5][
             0] > visibility_th:
-        cv.line(image, landmark_point[4][1], landmark_point[5][1],
+        cv2.line(image, landmark_point[4][1], landmark_point[5][1],
                 (0, 255, 0), 2)
     if landmark_point[5][0] > visibility_th and landmark_point[6][
             0] > visibility_th:
-        cv.line(image, landmark_point[5][1], landmark_point[6][1],
+        cv2.line(image, landmark_point[5][1], landmark_point[6][1],
                 (0, 255, 0), 2)
 
+
     if landmark_point[9][0] > visibility_th and landmark_point[10][
             0] > visibility_th:
-        cv.line(image, landmark_point[9][1], landmark_point[10][1],
+        cv2.line(image, landmark_point[9][1], landmark_point[10][1],
                 (0, 255, 0), 2)
 
+
     if landmark_point[11][0] > visibility_th and landmark_point[12][
             0] > visibility_th:
-        cv.line(image, landmark_point[11][1], landmark_point[12][1],
+        cv2.line(image, landmark_point[11][1], landmark_point[12][1],
                 (0, 255, 0), 2)
 
     if landmark_point[11][0] > visibility_th and landmark_point[13][
             0] > visibility_th:
-        cv.line(image, landmark_point[11][1], landmark_point[13][1],
+        cv2.line(image, landmark_point[11][1], landmark_point[13][1],
                 (0, 255, 0), 2)
     if landmark_point[13][0] > visibility_th and landmark_point[15][
             0] > visibility_th:
-        cv.line(image, landmark_point[13][1], landmark_point[15][1],
+        cv2.line(image, landmark_point[13][1], landmark_point[15][1],
                 (0, 255, 0), 2)
 
     if landmark_point[12][0] > visibility_th and landmark_point[14][
             0] > visibility_th:
-        cv.line(image, landmark_point[12][1], landmark_point[14][1],
+        cv2.line(image, landmark_point[12][1], landmark_point[14][1],
                 (0, 255, 0), 2)
     if landmark_point[14][0] > visibility_th and landmark_point[16][
             0] > visibility_th:
-        cv.line(image, landmark_point[14][1], landmark_point[16][1],
+        cv2.line(image, landmark_point[14][1], landmark_point[16][1],
                 (0, 255, 0), 2)
 
     if landmark_point[15][0] > visibility_th and landmark_point[17][
             0] > visibility_th:
-        cv.line(image, landmark_point[15][1], landmark_point[17][1],
+        cv2.line(image, landmark_point[15][1], landmark_point[17][1],
                 (0, 255, 0), 2)
     if landmark_point[17][0] > visibility_th and landmark_point[19][
             0] > visibility_th:
-        cv.line(image, landmark_point[17][1], landmark_point[19][1],
+        cv2.line(image, landmark_point[17][1], landmark_point[19][1],
                 (0, 255, 0), 2)
     if landmark_point[19][0] > visibility_th and landmark_point[21][
             0] > visibility_th:
-        cv.line(image, landmark_point[19][1], landmark_point[21][1],
+        cv2.line(image, landmark_point[19][1], landmark_point[21][1],
                 (0, 255, 0), 2)
     if landmark_point[21][0] > visibility_th and landmark_point[15][
             0] > visibility_th:
-        cv.line(image, landmark_point[21][1], landmark_point[15][1],
+        cv2.line(image, landmark_point[21][1], landmark_point[15][1],
                 (0, 255, 0), 2)
 
     if landmark_point[16][0] > visibility_th and landmark_point[18][
             0] > visibility_th:
-        cv.line(image, landmark_point[16][1], landmark_point[18][1],
+        cv2.line(image, landmark_point[16][1], landmark_point[18][1],
                 (0, 255, 0), 2)
     if landmark_point[18][0] > visibility_th and landmark_point[20][
             0] > visibility_th:
-        cv.line(image, landmark_point[18][1], landmark_point[20][1],
+        cv2.line(image, landmark_point[18][1], landmark_point[20][1],
                 (0, 255, 0), 2)
     if landmark_point[20][0] > visibility_th and landmark_point[22][
             0] > visibility_th:
-        cv.line(image, landmark_point[20][1], landmark_point[22][1],
+        cv2.line(image, landmark_point[20][1], landmark_point[22][1],
                 (0, 255, 0), 2)
     if landmark_point[22][0] > visibility_th and landmark_point[16][
             0] > visibility_th:
-        cv.line(image, landmark_point[22][1], landmark_point[16][1],
+        cv2.line(image, landmark_point[22][1], landmark_point[16][1],
                 (0, 255, 0), 2)
 
     if landmark_point[11][0] > visibility_th and landmark_point[23][
             0] > visibility_th:
-        cv.line(image, landmark_point[11][1], landmark_point[23][1],
+        cv2.line(image, landmark_point[11][1], landmark_point[23][1],
                 (0, 255, 0), 2)
     if landmark_point[12][0] > visibility_th and landmark_point[24][
             0] > visibility_th:
-        cv.line(image, landmark_point[12][1], landmark_point[24][1],
+        cv2.line(image, landmark_point[12][1], landmark_point[24][1],
                 (0, 255, 0), 2)
     if landmark_point[23][0] > visibility_th and landmark_point[24][
             0] > visibility_th:
-        cv.line(image, landmark_point[23][1], landmark_point[24][1],
+        cv2.line(image, landmark_point[23][1], landmark_point[24][1],
                 (0, 255, 0), 2)
 
     if len(landmark_point) > 25:
         if landmark_point[23][0] > visibility_th and landmark_point[25][
                 0] > visibility_th:
-            cv.line(image, landmark_point[23][1], landmark_point[25][1],
+            cv2.line(image, landmark_point[23][1], landmark_point[25][1],
                     (0, 255, 0), 2)
         if landmark_point[25][0] > visibility_th and landmark_point[27][
                 0] > visibility_th:
-            cv.line(image, landmark_point[25][1], landmark_point[27][1],
+            cv2.line(image, landmark_point[25][1], landmark_point[27][1],
                     (0, 255, 0), 2)
         if landmark_point[27][0] > visibility_th and landmark_point[29][
                 0] > visibility_th:
-            cv.line(image, landmark_point[27][1], landmark_point[29][1],
+            cv2.line(image, landmark_point[27][1], landmark_point[29][1],
                     (0, 255, 0), 2)
         if landmark_point[29][0] > visibility_th and landmark_point[31][
                 0] > visibility_th:
-            cv.line(image, landmark_point[29][1], landmark_point[31][1],
+            cv2.line(image, landmark_point[29][1], landmark_point[31][1],
                     (0, 255, 0), 2)
 
 
         if landmark_point[24][0] > visibility_th and landmark_point[26][
                 0] > visibility_th:
-            cv.line(image, landmark_point[24][1], landmark_point[26][1],
+            cv2.line(image, landmark_point[24][1], landmark_point[26][1],
                     (0, 255, 0), 2)
         if landmark_point[26][0] > visibility_th and landmark_point[28][
                 0] > visibility_th:
-            cv.line(image, landmark_point[26][1], landmark_point[28][1],
+            cv2.line(image, landmark_point[26][1], landmark_point[28][1],
                     (0, 255, 0), 2)
         if landmark_point[28][0] > visibility_th and landmark_point[30][
                 0] > visibility_th:
-            cv.line(image, landmark_point[28][1], landmark_point[30][1],
+            cv2.line(image, landmark_point[28][1], landmark_point[30][1],
                     (0, 255, 0), 2)
         if landmark_point[30][0] > visibility_th and landmark_point[32][
                 0] > visibility_th:
-            cv.line(image, landmark_point[30][1], landmark_point[32][1],
+            cv2.line(image, landmark_point[30][1], landmark_point[32][1],
                     (0, 255, 0), 2)
     return image