Skip to content

Commit

Permalink
feat: add transfer keypoint to coordinate not only a image in util.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@zgx949
zgx949 committed Apr 13, 2024
1 parent 5ee71dc commit a7a252e
Showing 1 changed file with 45 additions and 17 deletions.
62 changes: 45 additions & 17 deletions src/util.py
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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]]
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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]
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit a7a252e

Please sign in to comment.