def get_pose_from_img_file(img_path):
import glob
import os
img = cv2.imread(img_path)
res = get_pose(img, res_only=True)
threshold = 0.9
dofs = []
for i in range(len(res['boxes'])):
if res["scores"][i] > threshold:
dofs.append(res['dofs'] * 180 / np.pi)
print("{} - dofs {}".format(img_path, dofs))
print("++++++++++++++++++++++++++++++++++++++")
def get_pose_from_imgs_folder(folder_path):
import glob
import os
import tqdm
folder_path = os.path.join(folder_path, "*.jpg")
for img_path in tqdm.tqdm(glob.glob(folder_path)):
img = cv2.imread(img_path)
res = get_pose(img, res_only=True)
threshold = 0.9
dofs = []
for i in range(len(res['boxes'])):
if res["scores"][i] > threshold:
dofs.append(res['dofs'] * 180 / np.pi)
print("{} - dofs {}".format(img_path, dofs))
print("++++++++++++++++++++++++++++++++++++++")
def test_pose():
cap = cv2.VideoCapture(0)
while True:
ret, frame = cap.read()
if not ret:
break
# frame = cv2.flip(frame, 0)
key = cv2.waitKey(1) & 0xFF
time_0 = time.time()
frame = get_pose(frame)
logging.info("reference time: {}".format(time.time() - time_0))
print("-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-")
cv2.imshow("", frame)
if key == ord("q"):
break
def __getitem__(self, index):
img, target = None, None
env = self.env
with env.begin(write=False) as txn:
byteflow = txn.get(self.keys[index])
data = msgpack.loads(byteflow)
# load image
imgbuf = data[0]
buf = six.BytesIO()
buf.write(imgbuf)
buf.seek(0)
img = Image.open(buf).convert("RGB")
# load landmarks label
landmark_labels = data[4]
# load bbox
bbox_labels = np.asarray(data[2])
# apply augmentations that are provided from the parent class in creation order
for augmentation_method in self.augmentation_methods:
img, bbox_labels, landmark_labels = augmentation_method(
img, bbox_labels, landmark_labels
)
# create global intrinsics
(img_w, img_h) = img.size
global_intrinsics = np.array(
[[img_w + img_h, 0, img_w // 2], [0, img_w + img_h, img_h // 2], [0, 0, 1]]
)
projected_bbox_labels = []
pose_labels = []
img = np.array(img)
# get pose labels
for i in range(len(bbox_labels)):
bbox = bbox_labels[i]
lms = np.asarray(landmark_labels[i])
# black out faces that do not have pose annotation
if -1 in lms:
img[int(bbox[1]) : int(bbox[3]), int(bbox[0]) : int(bbox[2]), :] = 0
continue
# convert landmarks to bbox
bbox_lms = lms.copy()
bbox_lms[:, 0] -= bbox[0]
bbox_lms[:, 1] -= bbox[1]
# create bbox intrinsincs
w = int(bbox[2] - bbox[0])
h = int(bbox[3] - bbox[1])
bbox_intrinsics = np.array(
[[w + h, 0, w // 2], [0, w + h, h // 2], [0, 0, 1]]
)
# get pose between gt points and 3D reference
if len(bbox_lms) == 5:
P, pose = get_pose(self.threed_5_points, bbox_lms, bbox_intrinsics)
else:
P, pose = get_pose(self.threed_68_points, bbox_lms, bbox_intrinsics)
# convert to global image
pose_label = pose_bbox_to_full_image(pose, global_intrinsics, bbox)
# project points and get bbox
projected_lms, _ = plot_3d_landmark(
self.threed_68_points, pose_label, global_intrinsics
)
projected_bbox = expand_bbox_rectangle(
img_w, img_h, 1.1, 1.1, projected_lms, roll=pose_label[2]
)
pose_labels.append(pose_label)
projected_bbox_labels.append(projected_bbox)
pose_labels = np.asarray(pose_labels)
img = Image.fromarray(img)
if self.transform is not None:
img = self.transform(img)
target = {
"dofs": torch.from_numpy(pose_labels).float(),
"boxes": torch.from_numpy(np.asarray(projected_bbox_labels)).float(),
"labels": torch.ones((len(bbox_labels),), dtype=torch.int64),
}
return img, target