-
Notifications
You must be signed in to change notification settings - Fork 0
/
openpose.py
127 lines (105 loc) · 5.3 KB
/
openpose.py
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
# Openpose
# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
# 3rd Edited by ControlNet
# 4th Edited by ControlNet (added face and correct hands)
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import torch
import numpy as np
import util
from body import Body
from hand import Hand
from face import Face
from util import annotator_ckpts_path
body_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth"
hand_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth"
face_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth"
def draw_pose(pose, H, W, draw_body=True, draw_hand=True, draw_face=True):
bodies = pose['bodies']
candidate = bodies['candidate']
subset = bodies['subset']
canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
if draw_body:
canvas = util.draw_bodypose(canvas, candidate, subset)
if draw_hand:
hands = pose['hands']
canvas = util.draw_handpose(canvas, hands)
if draw_face:
faces = pose['faces']
canvas = util.draw_facepose(canvas, faces)
return canvas
class OpenposeDetectorBodyOnly:
def __init__(self):
body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
if not os.path.exists(body_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)
self.body_estimation = Body(body_modelpath)
def __call__(self, oriImg, return_is_index=False):
oriImg = oriImg[:, :, ::-1].copy()
H, W, C = oriImg.shape
with torch.no_grad():
candidate, subset = self.body_estimation(oriImg)
if candidate.ndim == 2 and candidate.shape[1] == 4:
candidate = candidate[:, :2]
candidate[:, 0] /= float(W)
candidate[:, 1] /= float(H)
bodies = dict(candidate=candidate.tolist(), subset=subset.tolist())
pose = dict(bodies=bodies)
if return_is_index:
return draw_pose(pose, H, W, draw_hand=False, draw_face=False), pose
else:
return draw_pose(pose, H, W, draw_hand=False, draw_face=False)
class OpenposeDetector:
def __init__(self):
body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
face_modelpath = os.path.join(annotator_ckpts_path, "facenet.pth")
if not os.path.exists(body_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)
if not os.path.exists(hand_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)
if not os.path.exists(face_modelpath):
from basicsr.utils.download_util import load_file_from_url
load_file_from_url(face_model_path, model_dir=annotator_ckpts_path)
self.body_estimation = Body(body_modelpath)
self.hand_estimation = Hand(hand_modelpath)
self.face_estimation = Face(face_modelpath)
def __call__(self, oriImg, hand_and_face=False, return_is_index=False):
oriImg = oriImg[:, :, ::-1].copy()
H, W, C = oriImg.shape
with torch.no_grad():
candidate, subset = self.body_estimation(oriImg)
hands = []
faces = []
if hand_and_face:
# Hand
hands_list = util.handDetect(candidate, subset, oriImg)
for x, y, w, is_left in hands_list:
peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :]).astype(np.float32)
if peaks.ndim == 2 and peaks.shape[1] == 2:
peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
hands.append(peaks.tolist())
# Face
faces_list = util.faceDetect(candidate, subset, oriImg)
for x, y, w in faces_list:
heatmaps = self.face_estimation(oriImg[y:y+w, x:x+w, :])
peaks = self.face_estimation.compute_peaks_from_heatmaps(heatmaps).astype(np.float32)
if peaks.ndim == 2 and peaks.shape[1] == 2:
peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
faces.append(peaks.tolist())
if candidate.ndim == 2 and candidate.shape[1] == 4:
candidate = candidate[:, :2]
candidate[:, 0] /= float(W)
candidate[:, 1] /= float(H)
bodies = dict(candidate=candidate.tolist(), subset=subset.tolist())
pose = dict(bodies=bodies, hands=hands, faces=faces)
if return_is_index:
return pose
else:
return draw_pose(pose, H, W)