|
|
|
|
|
|
| import os
|
|
|
| import cv2
|
| import torch
|
| import numpy as np
|
| from . import util
|
| from .wholebody import Wholebody, HWC3, resize_image
|
| from PIL import Image
|
|
|
| os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
|
|
| def convert_to_numpy(image):
|
| if isinstance(image, Image.Image):
|
| image = np.array(image)
|
| elif isinstance(image, torch.Tensor):
|
| image = image.detach().cpu().numpy()
|
| elif isinstance(image, np.ndarray):
|
| image = image.copy()
|
| else:
|
| raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
| return image
|
|
|
|
|
|
|
| def draw_pose(pose, H, W, use_hand=False, use_body=False, use_face=False):
|
| bodies = pose['bodies']
|
| faces = pose['faces']
|
| hands = pose['hands']
|
| candidate = bodies['candidate']
|
| subset = bodies['subset']
|
| canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
|
|
|
| if use_body:
|
| canvas = util.draw_bodypose(canvas, candidate, subset)
|
| if use_hand:
|
| canvas = util.draw_handpose(canvas, hands)
|
| if use_face:
|
| canvas = util.draw_facepose(canvas, faces)
|
|
|
| return canvas
|
|
|
|
|
| class PoseAnnotator:
|
| def __init__(self, cfg, device=None):
|
| onnx_det = cfg['DETECTION_MODEL']
|
| onnx_pose = cfg['POSE_MODEL']
|
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
|
| self.pose_estimation = Wholebody(onnx_det, onnx_pose, device=self.device)
|
| self.resize_size = cfg.get("RESIZE_SIZE", 1024)
|
| self.use_body = cfg.get('USE_BODY', True)
|
| self.use_face = cfg.get('USE_FACE', True)
|
| self.use_hand = cfg.get('USE_HAND', True)
|
|
|
| @torch.no_grad()
|
| @torch.inference_mode
|
| def forward(self, image):
|
| image = convert_to_numpy(image)
|
| input_image = HWC3(image[..., ::-1])
|
| return self.process(resize_image(input_image, self.resize_size), image.shape[:2])
|
|
|
| def process(self, ori_img, ori_shape):
|
| ori_h, ori_w = ori_shape
|
| ori_img = ori_img.copy()
|
| H, W, C = ori_img.shape
|
| with torch.no_grad():
|
| candidate, subset, det_result = self.pose_estimation(ori_img)
|
| nums, keys, locs = candidate.shape
|
| candidate[..., 0] /= float(W)
|
| candidate[..., 1] /= float(H)
|
| body = candidate[:, :18].copy()
|
| body = body.reshape(nums * 18, locs)
|
| score = subset[:, :18]
|
| for i in range(len(score)):
|
| for j in range(len(score[i])):
|
| if score[i][j] > 0.3:
|
| score[i][j] = int(18 * i + j)
|
| else:
|
| score[i][j] = -1
|
|
|
| un_visible = subset < 0.3
|
| candidate[un_visible] = -1
|
|
|
| foot = candidate[:, 18:24]
|
|
|
| faces = candidate[:, 24:92]
|
|
|
| hands = candidate[:, 92:113]
|
| hands = np.vstack([hands, candidate[:, 113:]])
|
|
|
| bodies = dict(candidate=body, subset=score)
|
| pose = dict(bodies=bodies, hands=hands, faces=faces)
|
|
|
| ret_data = {}
|
| if self.use_body:
|
| detected_map_body = draw_pose(pose, H, W, use_body=True)
|
| detected_map_body = cv2.resize(detected_map_body[..., ::-1], (ori_w, ori_h),
|
| interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
|
| ret_data["detected_map_body"] = detected_map_body
|
|
|
| if self.use_face:
|
| detected_map_face = draw_pose(pose, H, W, use_face=True)
|
| detected_map_face = cv2.resize(detected_map_face[..., ::-1], (ori_w, ori_h),
|
| interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
|
| ret_data["detected_map_face"] = detected_map_face
|
|
|
| if self.use_body and self.use_face:
|
| detected_map_bodyface = draw_pose(pose, H, W, use_body=True, use_face=True)
|
| detected_map_bodyface = cv2.resize(detected_map_bodyface[..., ::-1], (ori_w, ori_h),
|
| interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
|
| ret_data["detected_map_bodyface"] = detected_map_bodyface
|
|
|
| if self.use_hand and self.use_body and self.use_face:
|
| detected_map_handbodyface = draw_pose(pose, H, W, use_hand=True, use_body=True, use_face=True)
|
| detected_map_handbodyface = cv2.resize(detected_map_handbodyface[..., ::-1], (ori_w, ori_h),
|
| interpolation=cv2.INTER_LANCZOS4 if ori_h * ori_w > H * W else cv2.INTER_AREA)
|
| ret_data["detected_map_handbodyface"] = detected_map_handbodyface
|
|
|
|
|
| if det_result.shape[0] > 0:
|
| w_ratio, h_ratio = ori_w / W, ori_h / H
|
| det_result[..., ::2] *= h_ratio
|
| det_result[..., 1::2] *= w_ratio
|
| det_result = det_result.astype(np.int32)
|
| return ret_data, det_result
|
|
|
|
|
| class PoseBodyFaceAnnotator(PoseAnnotator):
|
| def __init__(self, cfg):
|
| super().__init__(cfg)
|
| self.use_body, self.use_face, self.use_hand = True, True, False
|
| @torch.no_grad()
|
| @torch.inference_mode
|
| def forward(self, image):
|
| ret_data, det_result = super().forward(image)
|
| return ret_data['detected_map_bodyface']
|
|
|
|
|
| class PoseBodyFaceVideoAnnotator(PoseBodyFaceAnnotator):
|
| def forward(self, frames):
|
| ret_frames = []
|
| for frame in frames:
|
| anno_frame = super().forward(np.array(frame))
|
| ret_frames.append(anno_frame)
|
| return ret_frames
|
|
|
| import imageio
|
|
|
| def save_one_video(file_path, videos, fps=8, quality=8, macro_block_size=None):
|
| try:
|
| video_writer = imageio.get_writer(file_path, fps=fps, codec='libx264', quality=quality, macro_block_size=macro_block_size)
|
| for frame in videos:
|
| video_writer.append_data(frame)
|
| video_writer.close()
|
| return True
|
| except Exception as e:
|
| print(f"Video save error: {e}")
|
| return False
|
|
|
| def get_frames(video_path):
|
| frames = []
|
|
|
|
|
|
|
| cap = cv2.VideoCapture(video_path)
|
|
|
| fps = cap.get(cv2.CAP_PROP_FPS)
|
| print("video fps: " + str(fps))
|
| i = 0
|
| while cap.isOpened():
|
| ret, frame = cap.read()
|
| if ret == False:
|
| break
|
| frames.append(frame)
|
| i += 1
|
|
|
| cap.release()
|
| cv2.destroyAllWindows()
|
|
|
| return frames, fps
|
|
|
|
|
