feat: add transfer keypoint to coordinate not only a image in util.py

Author: zgx949

src/util.py

@@ -1,43 +1,66 @@
-import numpy as np
 import math
+
 import cv2
 import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 from matplotlib.figure import Figure
-import numpy as np
-import matplotlib.pyplot as plt
-import cv2
 
 
 def padRightDownCorner(img, stride, padValue):
     h = img.shape[0]
     w = img.shape[1]
 
     pad = 4 * [None]
-    pad[0] = 0 # up
-    pad[1] = 0 # left
-    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
-    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
+    pad[0] = 0  # up
+    pad[1] = 0  # left
+    pad[2] = 0 if (h % stride == 0) else stride - (h % stride)  # down
+    pad[3] = 0 if (w % stride == 0) else stride - (w % stride)  # right
 
     img_padded = img
-    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
+    pad_up = np.tile(img_padded[0:1, :, :] * 0 + padValue, (pad[0], 1, 1))
     img_padded = np.concatenate((pad_up, img_padded), axis=0)
-    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
+    pad_left = np.tile(img_padded[:, 0:1, :] * 0 + padValue, (1, pad[1], 1))
     img_padded = np.concatenate((pad_left, img_padded), axis=1)
-    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
+    pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + padValue, (pad[2], 1, 1))
     img_padded = np.concatenate((img_padded, pad_down), axis=0)
-    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
+    pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + padValue, (1, pad[3], 1))
     img_padded = np.concatenate((img_padded, pad_right), axis=1)
 
     return img_padded, pad
 
+
 # transfer caffe model to pytorch which will match the layer name
 def transfer(model, model_weights):
     transfered_model_weights = {}
     for weights_name in model.state_dict().keys():
         transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
     return transfered_model_weights
 
+
+def transfer2coordinate(candidate, subset):
+    coordinates = []
+    keyMap = ['nose', 'neck',
+              'left_shoulder', 'left_elbow', 'left_wrist',
+              'right_shoulder', 'right_elbow', 'right_wrist',
+              'left_hip', 'left_knee', 'left_ankle',
+              'right_hip', 'right_knee', 'right_ankle',
+              'left_eyebrow_peak', 'left_eyebrow_tail',
+              'right_eyebrow_peak', 'right_eyebrow_tail'
+              ]
+    for n in range(len(subset)):
+        keypoint = {}
+        for i in range(18):
+            index = int(subset[n][i])
+            if index == -1:
+                continue
+            x, y = candidate[index][0:2]
+            key = keyMap[i]
+            keypoint[key] = (x, y)
+        coordinates.append(keypoint)
+    return coordinates
+
 # draw the body keypoint and lims
 def draw_bodypose(canvas, candidate, subset):
     stickwidth = 4
@@ -74,6 +97,7 @@ def draw_bodypose(canvas, candidate, subset):
     # plt.imshow(canvas[:, :, [2, 1, 0]])
     return canvas
 
+
 def draw_handpose(canvas, all_hand_peaks, show_number=False):
     edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
              [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
@@ -90,10 +114,10 @@ def draw_handpose(canvas, all_hand_peaks, show_number=False):
 
     for peaks in all_hand_peaks:
         for ie, e in enumerate(edges):
-            if np.sum(np.all(peaks[e], axis=1)==0)==0:
+            if np.sum(np.all(peaks[e], axis=1) == 0) == 0:
                 x1, y1 = peaks[e[0]]
                 x2, y2 = peaks[e[1]]
-                ax.plot([x1, x2], [y1, y2], color=matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0]))
+                ax.plot([x1, x2], [y1, y2], color=matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]))
 
         for i, keyponit in enumerate(peaks):
             x, y = keyponit
@@ -104,17 +128,19 @@ def draw_handpose(canvas, all_hand_peaks, show_number=False):
     canvas = np.fromstring(bg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
     return canvas
 
+
 # image drawed by opencv is not good.
 def draw_handpose_by_opencv(canvas, peaks, show_number=False):
     edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
              [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
     # cv2.rectangle(canvas, (x, y), (x+w, y+w), (0, 255, 0), 2, lineType=cv2.LINE_AA)
     # cv2.putText(canvas, 'left' if is_left else 'right', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
     for ie, e in enumerate(edges):
-        if np.sum(np.all(peaks[e], axis=1)==0)==0:
+        if np.sum(np.all(peaks[e], axis=1) == 0) == 0:
             x1, y1 = peaks[e[0]]
             x2, y2 = peaks[e[1]]
-            cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0])*255, thickness=2)
+            cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
+                     thickness=2)
 
     for i, keyponit in enumerate(peaks):
         x, y = keyponit
@@ -123,6 +149,7 @@ def draw_handpose_by_opencv(canvas, peaks, show_number=False):
             cv2.putText(canvas, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), lineType=cv2.LINE_AA)
     return canvas
 
+
 # detect hand according to body pose keypoints
 # please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
 def handDetect(candidate, subset, oriImg):
@@ -138,7 +165,7 @@ def handDetect(candidate, subset, oriImg):
         if not (has_left or has_right):
             continue
         hands = []
-        #left hand
+        # left hand
         if has_left:
             left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
             x1, y1 = candidate[left_shoulder_index][:2]
@@ -189,6 +216,7 @@ def handDetect(candidate, subset, oriImg):
     '''
     return detect_result
 
+
 # get max index of 2d array
 def npmax(array):
     arrayindex = array.argmax(1)