diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/utils/animal_landmark_runner.py b/src/utils/animal_landmark_runner.py new file mode 100644 index 0000000..963acbd --- /dev/null +++ b/src/utils/animal_landmark_runner.py @@ -0,0 +1,138 @@ +# coding: utf-8 + +""" +face detectoin and alignment using XPose +""" + +import os +import pickle +import torch +import numpy as np +from PIL import Image +from torchvision.ops import nms + +from .timer import Timer +from .rprint import rlog as log +from .helper import clean_state_dict + +from .dependencies.XPose import transforms as T +from .dependencies.XPose.models import build_model +from .dependencies.XPose.predefined_keypoints import * +from .dependencies.XPose.util import box_ops +from .dependencies.XPose.util.config import Config + + +class XPoseRunner(object): + def __init__(self, model_config_path, model_checkpoint_path, embeddings_cache_path=None, cpu_only=False, **kwargs): + self.device_id = kwargs.get("device_id", 0) + self.flag_use_half_precision = kwargs.get("flag_use_half_precision", True) + self.device = f"cuda:{self.device_id}" if not cpu_only else "cpu" + self.model = self.load_animal_model(model_config_path, model_checkpoint_path, self.device) + self.timer = Timer() + # Load cached embeddings if available + try: + with open(f'{embeddings_cache_path}_9.pkl', 'rb') as f: + self.ins_text_embeddings_9, self.kpt_text_embeddings_9 = pickle.load(f) + with open(f'{embeddings_cache_path}_68.pkl', 'rb') as f: + self.ins_text_embeddings_68, self.kpt_text_embeddings_68 = pickle.load(f) + print("Loaded cached embeddings from file.") + except Exception: + raise ValueError("Could not load clip embeddings from file, please check your file path.") + + def load_animal_model(self, model_config_path, model_checkpoint_path, device): + args = Config.fromfile(model_config_path) + args.device = device + model = build_model(args) + checkpoint = torch.load(model_checkpoint_path, map_location=lambda storage, loc: storage) + load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) + model.eval() + return model + + def load_image(self, input_image): + image_pil = input_image.convert("RGB") + transform = T.Compose([ + T.RandomResize([800], max_size=1333), # NOTE: fixed size to 800 + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), + ]) + image, _ = transform(image_pil, None) + return image_pil, image + + def get_unipose_output(self, image, instance_text_prompt, keypoint_text_prompt, box_threshold, IoU_threshold): + instance_list = instance_text_prompt.split(',') + + if len(keypoint_text_prompt) == 9: + # torch.Size([1, 512]) torch.Size([9, 512]) + ins_text_embeddings, kpt_text_embeddings = self.ins_text_embeddings_9, self.kpt_text_embeddings_9 + elif len(keypoint_text_prompt) ==68: + # torch.Size([1, 512]) torch.Size([68, 512]) + ins_text_embeddings, kpt_text_embeddings = self.ins_text_embeddings_68, self.kpt_text_embeddings_68 + else: + raise ValueError("Invalid number of keypoint embeddings.") + target = { + "instance_text_prompt": instance_list, + "keypoint_text_prompt": keypoint_text_prompt, + "object_embeddings_text": ins_text_embeddings.float(), + "kpts_embeddings_text": torch.cat((kpt_text_embeddings.float(), torch.zeros(100 - kpt_text_embeddings.shape[0], 512, device=self.device)), dim=0), + "kpt_vis_text": torch.cat((torch.ones(kpt_text_embeddings.shape[0], device=self.device), torch.zeros(100 - kpt_text_embeddings.shape[0], device=self.device)), dim=0) + } + + self.model = self.model.to(self.device) + image = image.to(self.device) + + with torch.no_grad(): + with torch.autocast(device_type=self.device[:4], dtype=torch.float16, enabled=self.flag_use_half_precision): + outputs = self.model(image[None], [target]) + + logits = outputs["pred_logits"].sigmoid()[0] + boxes = outputs["pred_boxes"][0] + keypoints = outputs["pred_keypoints"][0][:, :2 * len(keypoint_text_prompt)] + + logits_filt = logits.cpu().clone() + boxes_filt = boxes.cpu().clone() + keypoints_filt = keypoints.cpu().clone() + filt_mask = logits_filt.max(dim=1)[0] > box_threshold + logits_filt = logits_filt[filt_mask] + boxes_filt = boxes_filt[filt_mask] + keypoints_filt = keypoints_filt[filt_mask] + + keep_indices = nms(box_ops.box_cxcywh_to_xyxy(boxes_filt), logits_filt.max(dim=1)[0], iou_threshold=IoU_threshold) + + filtered_boxes = boxes_filt[keep_indices] + filtered_keypoints = keypoints_filt[keep_indices] + + return filtered_boxes, filtered_keypoints + + def run(self, input_image, instance_text_prompt, keypoint_text_example, box_threshold, IoU_threshold): + if keypoint_text_example in globals(): + keypoint_dict = globals()[keypoint_text_example] + elif instance_text_prompt in globals(): + keypoint_dict = globals()[instance_text_prompt] + else: + keypoint_dict = globals()["animal"] + + keypoint_text_prompt = keypoint_dict.get("keypoints") + keypoint_skeleton = keypoint_dict.get("skeleton") + + image_pil, image = self.load_image(input_image) + boxes_filt, keypoints_filt = self.get_unipose_output(image, instance_text_prompt, keypoint_text_prompt, box_threshold, IoU_threshold) + + size = image_pil.size + H, W = size[1], size[0] + keypoints_filt = keypoints_filt[0].squeeze(0) + kp = np.array(keypoints_filt.cpu()) + num_kpts = len(keypoint_text_prompt) + Z = kp[:num_kpts * 2] * np.array([W, H] * num_kpts) + Z = Z.reshape(num_kpts * 2) + x = Z[0::2] + y = Z[1::2] + return np.stack((x, y), axis=1) + + def warmup(self): + self.timer.tic() + + img_rgb = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)) + self.run(img_rgb, 'face', 'face', box_threshold=0.0, IoU_threshold=0.0) + + elapse = self.timer.toc() + log(f'XPoseRunner warmup time: {elapse:.3f}s') diff --git a/src/utils/camera.py b/src/utils/camera.py new file mode 100644 index 0000000..485a1f7 --- /dev/null +++ b/src/utils/camera.py @@ -0,0 +1,73 @@ +# coding: utf-8 + +""" +functions for processing and transforming 3D facial keypoints +""" + +import numpy as np +import torch +import torch.nn.functional as F + +PI = np.pi + + +def headpose_pred_to_degree(pred): + """ + pred: (bs, 66) or (bs, 1) or others + """ + if pred.ndim > 1 and pred.shape[1] == 66: + # NOTE: note that the average is modified to 97.5 + device = pred.device + idx_tensor = [idx for idx in range(0, 66)] + idx_tensor = torch.FloatTensor(idx_tensor).to(device) + pred = F.softmax(pred, dim=1) + degree = torch.sum(pred*idx_tensor, axis=1) * 3 - 97.5 + + return degree + + return pred + + +def get_rotation_matrix(pitch_, yaw_, roll_): + """ the input is in degree + """ + # transform to radian + pitch = pitch_ / 180 * PI + yaw = yaw_ / 180 * PI + roll = roll_ / 180 * PI + + device = pitch.device + + if pitch.ndim == 1: + pitch = pitch.unsqueeze(1) + if yaw.ndim == 1: + yaw = yaw.unsqueeze(1) + if roll.ndim == 1: + roll = roll.unsqueeze(1) + + # calculate the euler matrix + bs = pitch.shape[0] + ones = torch.ones([bs, 1]).to(device) + zeros = torch.zeros([bs, 1]).to(device) + x, y, z = pitch, yaw, roll + + rot_x = torch.cat([ + ones, zeros, zeros, + zeros, torch.cos(x), -torch.sin(x), + zeros, torch.sin(x), torch.cos(x) + ], dim=1).reshape([bs, 3, 3]) + + rot_y = torch.cat([ + torch.cos(y), zeros, torch.sin(y), + zeros, ones, zeros, + -torch.sin(y), zeros, torch.cos(y) + ], dim=1).reshape([bs, 3, 3]) + + rot_z = torch.cat([ + torch.cos(z), -torch.sin(z), zeros, + torch.sin(z), torch.cos(z), zeros, + zeros, zeros, ones + ], dim=1).reshape([bs, 3, 3]) + + rot = rot_z @ rot_y @ rot_x + return rot.permute(0, 2, 1) # transpose diff --git a/src/utils/check_windows_port.py b/src/utils/check_windows_port.py new file mode 100644 index 0000000..2a5f928 --- /dev/null +++ b/src/utils/check_windows_port.py @@ -0,0 +1,18 @@ +import socket +import sys + +if len(sys.argv) != 2: + print("Usage: python check_port.py ") + sys.exit(1) + +port = int(sys.argv[1]) + +sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) +sock.settimeout(1) +result = sock.connect_ex(('127.0.0.1', port)) + +if result == 0: + print("LISTENING") +else: + print("NOT LISTENING") +sock.close diff --git a/src/utils/crop.py b/src/utils/crop.py new file mode 100644 index 0000000..5671e95 --- /dev/null +++ b/src/utils/crop.py @@ -0,0 +1,423 @@ +# coding: utf-8 + +""" +cropping function and the related preprocess functions for cropping +""" + +import numpy as np +import os.path as osp +from math import sin, cos, acos, degrees +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) # NOTE: enforce single thread +from .rprint import rprint as print + +DTYPE = np.float32 +CV2_INTERP = cv2.INTER_LINEAR + +def make_abs_path(fn): + return osp.join(osp.dirname(osp.realpath(__file__)), fn) + +def _transform_img(img, M, dsize, flags=CV2_INTERP, borderMode=None): + """ conduct similarity or affine transformation to the image, do not do border operation! + img: + M: 2x3 matrix or 3x3 matrix + dsize: target shape (width, height) + """ + if isinstance(dsize, tuple) or isinstance(dsize, list): + _dsize = tuple(dsize) + else: + _dsize = (dsize, dsize) + + if borderMode is not None: + return cv2.warpAffine(img, M[:2, :], dsize=_dsize, flags=flags, borderMode=borderMode, borderValue=(0, 0, 0)) + else: + return cv2.warpAffine(img, M[:2, :], dsize=_dsize, flags=flags) + + +def _transform_pts(pts, M): + """ conduct similarity or affine transformation to the pts + pts: Nx2 ndarray + M: 2x3 matrix or 3x3 matrix + return: Nx2 + """ + return pts @ M[:2, :2].T + M[:2, 2] + + +def parse_pt2_from_pt101(pt101, use_lip=True): + """ + parsing the 2 points according to the 101 points, which cancels the roll + """ + # the former version use the eye center, but it is not robust, now use interpolation + pt_left_eye = np.mean(pt101[[39, 42, 45, 48]], axis=0) # left eye center + pt_right_eye = np.mean(pt101[[51, 54, 57, 60]], axis=0) # right eye center + + if use_lip: + # use lip + pt_center_eye = (pt_left_eye + pt_right_eye) / 2 + pt_center_lip = (pt101[75] + pt101[81]) / 2 + pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0) + else: + pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0) + return pt2 + + +def parse_pt2_from_pt106(pt106, use_lip=True): + """ + parsing the 2 points according to the 106 points, which cancels the roll + """ + pt_left_eye = np.mean(pt106[[33, 35, 40, 39]], axis=0) # left eye center + pt_right_eye = np.mean(pt106[[87, 89, 94, 93]], axis=0) # right eye center + + if use_lip: + # use lip + pt_center_eye = (pt_left_eye + pt_right_eye) / 2 + pt_center_lip = (pt106[52] + pt106[61]) / 2 + pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0) + else: + pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0) + return pt2 + + +def parse_pt2_from_pt203(pt203, use_lip=True): + """ + parsing the 2 points according to the 203 points, which cancels the roll + """ + pt_left_eye = np.mean(pt203[[0, 6, 12, 18]], axis=0) # left eye center + pt_right_eye = np.mean(pt203[[24, 30, 36, 42]], axis=0) # right eye center + if use_lip: + # use lip + pt_center_eye = (pt_left_eye + pt_right_eye) / 2 + pt_center_lip = (pt203[48] + pt203[66]) / 2 + pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0) + else: + pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0) + return pt2 + + +def parse_pt2_from_pt68(pt68, use_lip=True): + """ + parsing the 2 points according to the 68 points, which cancels the roll + """ + lm_idx = np.array([31, 37, 40, 43, 46, 49, 55], dtype=np.int32) - 1 + if use_lip: + pt5 = np.stack([ + np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye + np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye + pt68[lm_idx[0], :], # nose + pt68[lm_idx[5], :], # lip + pt68[lm_idx[6], :] # lip + ], axis=0) + + pt2 = np.stack([ + (pt5[0] + pt5[1]) / 2, + (pt5[3] + pt5[4]) / 2 + ], axis=0) + else: + pt2 = np.stack([ + np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye + np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye + ], axis=0) + + return pt2 + + +def parse_pt2_from_pt5(pt5, use_lip=True): + """ + parsing the 2 points according to the 5 points, which cancels the roll + """ + if use_lip: + pt2 = np.stack([ + (pt5[0] + pt5[1]) / 2, + (pt5[3] + pt5[4]) / 2 + ], axis=0) + else: + pt2 = np.stack([ + pt5[0], + pt5[1] + ], axis=0) + return pt2 + +def parse_pt2_from_pt9(pt9, use_lip=True): + ''' + parsing the 2 points according to the 9 points, which cancels the roll + ['right eye right', 'right eye left', 'left eye right', 'left eye left', 'nose tip', 'lip right', 'lip left', 'upper lip', 'lower lip'] + ''' + if use_lip: + pt9 = np.stack([ + (pt9[2] + pt9[3]) / 2, # left eye + (pt9[0] + pt9[1]) / 2, # right eye + pt9[4], + (pt9[5] + pt9[6] ) / 2 # lip + ], axis=0) + pt2 = np.stack([ + (pt9[0] + pt9[1]) / 2, # eye + pt9[3] # lip + ], axis=0) + else: + pt2 = np.stack([ + (pt9[2] + pt9[3]) / 2, + (pt9[0] + pt9[1]) / 2, + ], axis=0) + + return pt2 + +def parse_pt2_from_pt_x(pts, use_lip=True): + if pts.shape[0] == 101: + pt2 = parse_pt2_from_pt101(pts, use_lip=use_lip) + elif pts.shape[0] == 106: + pt2 = parse_pt2_from_pt106(pts, use_lip=use_lip) + elif pts.shape[0] == 68: + pt2 = parse_pt2_from_pt68(pts, use_lip=use_lip) + elif pts.shape[0] == 5: + pt2 = parse_pt2_from_pt5(pts, use_lip=use_lip) + elif pts.shape[0] == 203: + pt2 = parse_pt2_from_pt203(pts, use_lip=use_lip) + elif pts.shape[0] > 101: + # take the first 101 points + pt2 = parse_pt2_from_pt101(pts[:101], use_lip=use_lip) + elif pts.shape[0] == 9: + pt2 = parse_pt2_from_pt9(pts, use_lip=use_lip) + else: + raise Exception(f'Unknow shape: {pts.shape}') + + if not use_lip: + # NOTE: to compile with the latter code, need to rotate the pt2 90 degrees clockwise manually + v = pt2[1] - pt2[0] + pt2[1, 0] = pt2[0, 0] - v[1] + pt2[1, 1] = pt2[0, 1] + v[0] + + return pt2 + + +def parse_rect_from_landmark( + pts, + scale=1.5, + need_square=True, + vx_ratio=0, + vy_ratio=0, + use_deg_flag=False, + **kwargs +): + """parsing center, size, angle from 101/68/5/x landmarks + vx_ratio: the offset ratio along the pupil axis x-axis, multiplied by size + vy_ratio: the offset ratio along the pupil axis y-axis, multiplied by size, which is used to contain more forehead area + + judge with pts.shape + """ + pt2 = parse_pt2_from_pt_x(pts, use_lip=kwargs.get('use_lip', True)) + + uy = pt2[1] - pt2[0] + l = np.linalg.norm(uy) + if l <= 1e-3: + uy = np.array([0, 1], dtype=DTYPE) + else: + uy /= l + ux = np.array((uy[1], -uy[0]), dtype=DTYPE) + + # the rotation degree of the x-axis, the clockwise is positive, the counterclockwise is negative (image coordinate system) + # print(uy) + # print(ux) + angle = acos(ux[0]) + if ux[1] < 0: + angle = -angle + + # rotation matrix + M = np.array([ux, uy]) + + # calculate the size which contains the angle degree of the bbox, and the center + center0 = np.mean(pts, axis=0) + rpts = (pts - center0) @ M.T # (M @ P.T).T = P @ M.T + lt_pt = np.min(rpts, axis=0) + rb_pt = np.max(rpts, axis=0) + center1 = (lt_pt + rb_pt) / 2 + + size = rb_pt - lt_pt + if need_square: + m = max(size[0], size[1]) + size[0] = m + size[1] = m + + size *= scale # scale size + center = center0 + ux * center1[0] + uy * center1[1] # counterclockwise rotation, equivalent to M.T @ center1.T + center = center + ux * (vx_ratio * size) + uy * \ + (vy_ratio * size) # considering the offset in vx and vy direction + + if use_deg_flag: + angle = degrees(angle) + + return center, size, angle + + +def parse_bbox_from_landmark(pts, **kwargs): + center, size, angle = parse_rect_from_landmark(pts, **kwargs) + cx, cy = center + w, h = size + + # calculate the vertex positions before rotation + bbox = np.array([ + [cx-w/2, cy-h/2], # left, top + [cx+w/2, cy-h/2], + [cx+w/2, cy+h/2], # right, bottom + [cx-w/2, cy+h/2] + ], dtype=DTYPE) + + # construct rotation matrix + bbox_rot = bbox.copy() + R = np.array([ + [np.cos(angle), -np.sin(angle)], + [np.sin(angle), np.cos(angle)] + ], dtype=DTYPE) + + # calculate the relative position of each vertex from the rotation center, then rotate these positions, and finally add the coordinates of the rotation center + bbox_rot = (bbox_rot - center) @ R.T + center + + return { + 'center': center, # 2x1 + 'size': size, # scalar + 'angle': angle, # rad, counterclockwise + 'bbox': bbox, # 4x2 + 'bbox_rot': bbox_rot, # 4x2 + } + + +def crop_image_by_bbox(img, bbox, lmk=None, dsize=512, angle=None, flag_rot=False, **kwargs): + left, top, right, bot = bbox + if int(right - left) != int(bot - top): + print(f'right-left {right-left} != bot-top {bot-top}') + size = right - left + + src_center = np.array([(left + right) / 2, (top + bot) / 2], dtype=DTYPE) + tgt_center = np.array([dsize / 2, dsize / 2], dtype=DTYPE) + + s = dsize / size # scale + if flag_rot and angle is not None: + costheta, sintheta = cos(angle), sin(angle) + cx, cy = src_center[0], src_center[1] # ori center + tcx, tcy = tgt_center[0], tgt_center[1] # target center + # need to infer + M_o2c = np.array( + [[s * costheta, s * sintheta, tcx - s * (costheta * cx + sintheta * cy)], + [-s * sintheta, s * costheta, tcy - s * (-sintheta * cx + costheta * cy)]], + dtype=DTYPE + ) + else: + M_o2c = np.array( + [[s, 0, tgt_center[0] - s * src_center[0]], + [0, s, tgt_center[1] - s * src_center[1]]], + dtype=DTYPE + ) + + # if flag_rot and angle is None: + # print('angle is None, but flag_rotate is True', style="bold yellow") + + img_crop = _transform_img(img, M_o2c, dsize=dsize, borderMode=kwargs.get('borderMode', None)) + lmk_crop = _transform_pts(lmk, M_o2c) if lmk is not None else None + + M_o2c = np.vstack([M_o2c, np.array([0, 0, 1], dtype=DTYPE)]) + M_c2o = np.linalg.inv(M_o2c) + + # cv2.imwrite('crop.jpg', img_crop) + + return { + 'img_crop': img_crop, + 'lmk_crop': lmk_crop, + 'M_o2c': M_o2c, + 'M_c2o': M_c2o, + } + + +def _estimate_similar_transform_from_pts( + pts, + dsize, + scale=1.5, + vx_ratio=0, + vy_ratio=-0.1, + flag_do_rot=True, + **kwargs +): + """ calculate the affine matrix of the cropped image from sparse points, the original image to the cropped image, the inverse is the cropped image to the original image + pts: landmark, 101 or 68 points or other points, Nx2 + scale: the larger scale factor, the smaller face ratio + vx_ratio: x shift + vy_ratio: y shift, the smaller the y shift, the lower the face region + rot_flag: if it is true, conduct correction + """ + center, size, angle = parse_rect_from_landmark( + pts, scale=scale, vx_ratio=vx_ratio, vy_ratio=vy_ratio, + use_lip=kwargs.get('use_lip', True) + ) + + s = dsize / size[0] # scale + tgt_center = np.array([dsize / 2, dsize / 2], dtype=DTYPE) # center of dsize + + if flag_do_rot: + costheta, sintheta = cos(angle), sin(angle) + cx, cy = center[0], center[1] # ori center + tcx, tcy = tgt_center[0], tgt_center[1] # target center + # need to infer + M_INV = np.array( + [[s * costheta, s * sintheta, tcx - s * (costheta * cx + sintheta * cy)], + [-s * sintheta, s * costheta, tcy - s * (-sintheta * cx + costheta * cy)]], + dtype=DTYPE + ) + else: + M_INV = np.array( + [[s, 0, tgt_center[0] - s * center[0]], + [0, s, tgt_center[1] - s * center[1]]], + dtype=DTYPE + ) + + M_INV_H = np.vstack([M_INV, np.array([0, 0, 1])]) + M = np.linalg.inv(M_INV_H) + + # M_INV is from the original image to the cropped image, M is from the cropped image to the original image + return M_INV, M[:2, ...] + + +def crop_image(img, pts: np.ndarray, **kwargs): + dsize = kwargs.get('dsize', 224) + scale = kwargs.get('scale', 1.5) # 1.5 | 1.6 + vy_ratio = kwargs.get('vy_ratio', -0.1) # -0.0625 | -0.1 + + M_INV, _ = _estimate_similar_transform_from_pts( + pts, + dsize=dsize, + scale=scale, + vy_ratio=vy_ratio, + flag_do_rot=kwargs.get('flag_do_rot', True), + ) + + img_crop = _transform_img(img, M_INV, dsize) # origin to crop + pt_crop = _transform_pts(pts, M_INV) + + M_o2c = np.vstack([M_INV, np.array([0, 0, 1], dtype=DTYPE)]) + M_c2o = np.linalg.inv(M_o2c) + + ret_dct = { + 'M_o2c': M_o2c, # from the original image to the cropped image 3x3 + 'M_c2o': M_c2o, # from the cropped image to the original image 3x3 + 'img_crop': img_crop, # the cropped image + 'pt_crop': pt_crop, # the landmarks of the cropped image + } + + return ret_dct + +def average_bbox_lst(bbox_lst): + if len(bbox_lst) == 0: + return None + bbox_arr = np.array(bbox_lst) + return np.mean(bbox_arr, axis=0).tolist() + +def prepare_paste_back(mask_crop, crop_M_c2o, dsize): + """prepare mask for later image paste back + """ + mask_ori = _transform_img(mask_crop, crop_M_c2o, dsize) + mask_ori = mask_ori.astype(np.float32) / 255. + return mask_ori + +def paste_back(img_crop, M_c2o, img_ori, mask_ori): + """paste back the image + """ + dsize = (img_ori.shape[1], img_ori.shape[0]) + result = _transform_img(img_crop, M_c2o, dsize=dsize) + result = np.clip(mask_ori * result + (1 - mask_ori) * img_ori, 0, 255).astype(np.uint8) + return result diff --git a/src/utils/cropper.py b/src/utils/cropper.py new file mode 100644 index 0000000..f07bc8c --- /dev/null +++ b/src/utils/cropper.py @@ -0,0 +1,313 @@ +# coding: utf-8 + +import os.path as osp +import torch +import numpy as np +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) + +from PIL import Image +from typing import List, Tuple, Union +from dataclasses import dataclass, field + +from ..config.crop_config import CropConfig +from .crop import ( + average_bbox_lst, + crop_image, + crop_image_by_bbox, + parse_bbox_from_landmark, +) +from .io import contiguous +from .rprint import rlog as log +from .face_analysis_diy import FaceAnalysisDIY +from .human_landmark_runner import LandmarkRunner as HumanLandmark + +def make_abs_path(fn): + return osp.join(osp.dirname(osp.realpath(__file__)), fn) + + +@dataclass +class Trajectory: + start: int = -1 # start frame + end: int = -1 # end frame + lmk_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # lmk list + bbox_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # bbox list + M_c2o_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # M_c2o list + + frame_rgb_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame list + lmk_crop_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # lmk list + frame_rgb_crop_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame crop list + + +class Cropper(object): + def __init__(self, **kwargs) -> None: + self.crop_cfg: CropConfig = kwargs.get("crop_cfg", None) + self.image_type = kwargs.get("image_type", 'human_face') + device_id = kwargs.get("device_id", 0) + flag_force_cpu = kwargs.get("flag_force_cpu", False) + if flag_force_cpu: + device = "cpu" + face_analysis_wrapper_provider = ["CPUExecutionProvider"] + else: + try: + if torch.backends.mps.is_available(): + # Shape inference currently fails with CoreMLExecutionProvider + # for the retinaface model + device = "mps" + face_analysis_wrapper_provider = ["CPUExecutionProvider"] + else: + device = "cuda" + face_analysis_wrapper_provider = ["CUDAExecutionProvider"] + except: + device = "cuda" + face_analysis_wrapper_provider = ["CUDAExecutionProvider"] + self.face_analysis_wrapper = FaceAnalysisDIY( + name="buffalo_l", + root=self.crop_cfg.insightface_root, + providers=face_analysis_wrapper_provider, + ) + self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512), det_thresh=self.crop_cfg.det_thresh) + self.face_analysis_wrapper.warmup() + + self.human_landmark_runner = HumanLandmark( + ckpt_path=self.crop_cfg.landmark_ckpt_path, + onnx_provider=device, + device_id=device_id, + ) + self.human_landmark_runner.warmup() + + if self.image_type == "animal_face": + from .animal_landmark_runner import XPoseRunner as AnimalLandmarkRunner + self.animal_landmark_runner = AnimalLandmarkRunner( + model_config_path=self.crop_cfg.xpose_config_file_path, + model_checkpoint_path=self.crop_cfg.xpose_ckpt_path, + embeddings_cache_path=self.crop_cfg.xpose_embedding_cache_path, + flag_use_half_precision=kwargs.get("flag_use_half_precision", True), + ) + self.animal_landmark_runner.warmup() + + def update_config(self, user_args): + for k, v in user_args.items(): + if hasattr(self.crop_cfg, k): + setattr(self.crop_cfg, k, v) + + def crop_source_image(self, img_rgb_: np.ndarray, crop_cfg: CropConfig): + # crop a source image and get neccessary information + img_rgb = img_rgb_.copy() # copy it + img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) + + if self.image_type == "human_face": + src_face = self.face_analysis_wrapper.get( + img_bgr, + flag_do_landmark_2d_106=True, + direction=crop_cfg.direction, + max_face_num=crop_cfg.max_face_num, + ) + + if len(src_face) == 0: + log("No face detected in the source image.") + return None + elif len(src_face) > 1: + log(f"More than one face detected in the image, only pick one face by rule {crop_cfg.direction}.") + + # NOTE: temporarily only pick the first face, to support multiple face in the future + src_face = src_face[0] + lmk = src_face.landmark_2d_106 # this is the 106 landmarks from insightface + else: + tmp_dct = { + 'animal_face_9': 'animal_face', + 'animal_face_68': 'face' + } + + img_rgb_pil = Image.fromarray(img_rgb) + lmk = self.animal_landmark_runner.run( + img_rgb_pil, + 'face', + tmp_dct[crop_cfg.animal_face_type], + 0, + 0 + ) + + # crop the face + ret_dct = crop_image( + img_rgb, # ndarray + lmk, # 106x2 or Nx2 + dsize=crop_cfg.dsize, + scale=crop_cfg.scale, + vx_ratio=crop_cfg.vx_ratio, + vy_ratio=crop_cfg.vy_ratio, + flag_do_rot=crop_cfg.flag_do_rot, + ) + + # update a 256x256 version for network input + ret_dct["img_crop_256x256"] = cv2.resize(ret_dct["img_crop"], (256, 256), interpolation=cv2.INTER_AREA) + if self.image_type == "human_face": + lmk = self.human_landmark_runner.run(img_rgb, lmk) + ret_dct["lmk_crop"] = lmk + ret_dct["lmk_crop_256x256"] = ret_dct["lmk_crop"] * 256 / crop_cfg.dsize + else: + # 68x2 or 9x2 + ret_dct["lmk_crop"] = lmk + + return ret_dct + + def calc_lmk_from_cropped_image(self, img_rgb_, **kwargs): + direction = kwargs.get("direction", "large-small") + src_face = self.face_analysis_wrapper.get( + contiguous(img_rgb_[..., ::-1]), # convert to BGR + flag_do_landmark_2d_106=True, + direction=direction, + ) + if len(src_face) == 0: + log("No face detected in the source image.") + return None + elif len(src_face) > 1: + log(f"More than one face detected in the image, only pick one face by rule {direction}.") + src_face = src_face[0] + lmk = src_face.landmark_2d_106 + lmk = self.human_landmark_runner.run(img_rgb_, lmk) + + return lmk + + # TODO: support skipping frame with NO FACE + def crop_source_video(self, source_rgb_lst, crop_cfg: CropConfig, **kwargs): + """Tracking based landmarks/alignment and cropping""" + trajectory = Trajectory() + direction = kwargs.get("direction", "large-small") + for idx, frame_rgb in enumerate(source_rgb_lst): + if idx == 0 or trajectory.start == -1: + src_face = self.face_analysis_wrapper.get( + contiguous(frame_rgb[..., ::-1]), + flag_do_landmark_2d_106=True, + direction=crop_cfg.direction, + max_face_num=crop_cfg.max_face_num, + ) + if len(src_face) == 0: + log(f"No face detected in the frame #{idx}") + continue + elif len(src_face) > 1: + log(f"More than one face detected in the source frame_{idx}, only pick one face by rule {direction}.") + src_face = src_face[0] + lmk = src_face.landmark_2d_106 + lmk = self.human_landmark_runner.run(frame_rgb, lmk) + trajectory.start, trajectory.end = idx, idx + else: + # TODO: add IOU check for tracking + lmk = self.human_landmark_runner.run(frame_rgb, trajectory.lmk_lst[-1]) + trajectory.end = idx + + trajectory.lmk_lst.append(lmk) + + # crop the face + ret_dct = crop_image( + frame_rgb, # ndarray + lmk, # 106x2 or Nx2 + dsize=crop_cfg.dsize, + scale=crop_cfg.scale, + vx_ratio=crop_cfg.vx_ratio, + vy_ratio=crop_cfg.vy_ratio, + flag_do_rot=crop_cfg.flag_do_rot, + ) + lmk = self.human_landmark_runner.run(frame_rgb, lmk) + ret_dct["lmk_crop"] = lmk + + # update a 256x256 version for network input + ret_dct["img_crop_256x256"] = cv2.resize(ret_dct["img_crop"], (256, 256), interpolation=cv2.INTER_AREA) + ret_dct["lmk_crop_256x256"] = ret_dct["lmk_crop"] * 256 / crop_cfg.dsize + + trajectory.frame_rgb_crop_lst.append(ret_dct["img_crop_256x256"]) + trajectory.lmk_crop_lst.append(ret_dct["lmk_crop_256x256"]) + trajectory.M_c2o_lst.append(ret_dct['M_c2o']) + + return { + "frame_crop_lst": trajectory.frame_rgb_crop_lst, + "lmk_crop_lst": trajectory.lmk_crop_lst, + "M_c2o_lst": trajectory.M_c2o_lst, + } + + def crop_driving_video(self, driving_rgb_lst, **kwargs): + """Tracking based landmarks/alignment and cropping""" + trajectory = Trajectory() + direction = kwargs.get("direction", "large-small") + for idx, frame_rgb in enumerate(driving_rgb_lst): + if idx == 0 or trajectory.start == -1: + src_face = self.face_analysis_wrapper.get( + contiguous(frame_rgb[..., ::-1]), + flag_do_landmark_2d_106=True, + direction=direction, + ) + if len(src_face) == 0: + log(f"No face detected in the frame #{idx}") + continue + elif len(src_face) > 1: + log(f"More than one face detected in the driving frame_{idx}, only pick one face by rule {direction}.") + src_face = src_face[0] + lmk = src_face.landmark_2d_106 + lmk = self.human_landmark_runner.run(frame_rgb, lmk) + trajectory.start, trajectory.end = idx, idx + else: + lmk = self.human_landmark_runner.run(frame_rgb, trajectory.lmk_lst[-1]) + trajectory.end = idx + + trajectory.lmk_lst.append(lmk) + ret_bbox = parse_bbox_from_landmark( + lmk, + scale=self.crop_cfg.scale_crop_driving_video, + vx_ratio_crop_driving_video=self.crop_cfg.vx_ratio_crop_driving_video, + vy_ratio=self.crop_cfg.vy_ratio_crop_driving_video, + )["bbox"] + bbox = [ + ret_bbox[0, 0], + ret_bbox[0, 1], + ret_bbox[2, 0], + ret_bbox[2, 1], + ] # 4, + trajectory.bbox_lst.append(bbox) # bbox + trajectory.frame_rgb_lst.append(frame_rgb) + + global_bbox = average_bbox_lst(trajectory.bbox_lst) + + for idx, (frame_rgb, lmk) in enumerate(zip(trajectory.frame_rgb_lst, trajectory.lmk_lst)): + ret_dct = crop_image_by_bbox( + frame_rgb, + global_bbox, + lmk=lmk, + dsize=kwargs.get("dsize", 512), + flag_rot=False, + borderValue=(0, 0, 0), + ) + trajectory.frame_rgb_crop_lst.append(ret_dct["img_crop"]) + trajectory.lmk_crop_lst.append(ret_dct["lmk_crop"]) + + return { + "frame_crop_lst": trajectory.frame_rgb_crop_lst, + "lmk_crop_lst": trajectory.lmk_crop_lst, + } + + + def calc_lmks_from_cropped_video(self, driving_rgb_crop_lst, **kwargs): + """Tracking based landmarks/alignment""" + trajectory = Trajectory() + direction = kwargs.get("direction", "large-small") + + for idx, frame_rgb_crop in enumerate(driving_rgb_crop_lst): + if idx == 0 or trajectory.start == -1: + src_face = self.face_analysis_wrapper.get( + contiguous(frame_rgb_crop[..., ::-1]), # convert to BGR + flag_do_landmark_2d_106=True, + direction=direction, + ) + if len(src_face) == 0: + log(f"No face detected in the frame #{idx}") + raise Exception(f"No face detected in the frame #{idx}") + elif len(src_face) > 1: + log(f"More than one face detected in the driving frame_{idx}, only pick one face by rule {direction}.") + src_face = src_face[0] + lmk = src_face.landmark_2d_106 + lmk = self.human_landmark_runner.run(frame_rgb_crop, lmk) + trajectory.start, trajectory.end = idx, idx + else: + lmk = self.human_landmark_runner.run(frame_rgb_crop, trajectory.lmk_lst[-1]) + trajectory.end = idx + + trajectory.lmk_lst.append(lmk) + return trajectory.lmk_lst diff --git a/src/utils/face_analysis_diy.py b/src/utils/face_analysis_diy.py new file mode 100644 index 0000000..359ac17 --- /dev/null +++ b/src/utils/face_analysis_diy.py @@ -0,0 +1,79 @@ +# coding: utf-8 + +""" +face detectoin and alignment using InsightFace +""" + +import numpy as np +from .rprint import rlog as log +from .dependencies.insightface.app import FaceAnalysis +from .dependencies.insightface.app.common import Face +from .timer import Timer + + +def sort_by_direction(faces, direction: str = 'large-small', face_center=None): + if len(faces) <= 0: + return faces + + if direction == 'left-right': + return sorted(faces, key=lambda face: face['bbox'][0]) + if direction == 'right-left': + return sorted(faces, key=lambda face: face['bbox'][0], reverse=True) + if direction == 'top-bottom': + return sorted(faces, key=lambda face: face['bbox'][1]) + if direction == 'bottom-top': + return sorted(faces, key=lambda face: face['bbox'][1], reverse=True) + if direction == 'small-large': + return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1])) + if direction == 'large-small': + return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]), reverse=True) + if direction == 'distance-from-retarget-face': + return sorted(faces, key=lambda face: (((face['bbox'][2]+face['bbox'][0])/2-face_center[0])**2+((face['bbox'][3]+face['bbox'][1])/2-face_center[1])**2)**0.5) + return faces + + +class FaceAnalysisDIY(FaceAnalysis): + def __init__(self, name='buffalo_l', root='~/.insightface', allowed_modules=None, **kwargs): + super().__init__(name=name, root=root, allowed_modules=allowed_modules, **kwargs) + + self.timer = Timer() + + def get(self, img_bgr, **kwargs): + max_num = kwargs.get('max_face_num', 0) # the number of the detected faces, 0 means no limit + flag_do_landmark_2d_106 = kwargs.get('flag_do_landmark_2d_106', True) # whether to do 106-point detection + direction = kwargs.get('direction', 'large-small') # sorting direction + face_center = None + + bboxes, kpss = self.det_model.detect(img_bgr, max_num=max_num, metric='default') + if bboxes.shape[0] == 0: + return [] + ret = [] + for i in range(bboxes.shape[0]): + bbox = bboxes[i, 0:4] + det_score = bboxes[i, 4] + kps = None + if kpss is not None: + kps = kpss[i] + face = Face(bbox=bbox, kps=kps, det_score=det_score) + for taskname, model in self.models.items(): + if taskname == 'detection': + continue + + if (not flag_do_landmark_2d_106) and taskname == 'landmark_2d_106': + continue + + # print(f'taskname: {taskname}') + model.get(img_bgr, face) + ret.append(face) + + ret = sort_by_direction(ret, direction, face_center) + return ret + + def warmup(self): + self.timer.tic() + + img_bgr = np.zeros((512, 512, 3), dtype=np.uint8) + self.get(img_bgr) + + elapse = self.timer.toc() + log(f'FaceAnalysisDIY warmup time: {elapse:.3f}s') diff --git a/src/utils/filter.py b/src/utils/filter.py new file mode 100644 index 0000000..a7b1692 --- /dev/null +++ b/src/utils/filter.py @@ -0,0 +1,19 @@ +# coding: utf-8 + +import torch +import numpy as np +from pykalman import KalmanFilter + + +def smooth(x_d_lst, shape, device, observation_variance=3e-7, process_variance=1e-5): + x_d_lst_reshape = [x.reshape(-1) for x in x_d_lst] + x_d_stacked = np.vstack(x_d_lst_reshape) + kf = KalmanFilter( + initial_state_mean=x_d_stacked[0], + n_dim_obs=x_d_stacked.shape[1], + transition_covariance=process_variance * np.eye(x_d_stacked.shape[1]), + observation_covariance=observation_variance * np.eye(x_d_stacked.shape[1]) + ) + smoothed_state_means, _ = kf.smooth(x_d_stacked) + x_d_lst_smooth = [torch.tensor(state_mean.reshape(shape[-2:]), dtype=torch.float32, device=device) for state_mean in smoothed_state_means] + return x_d_lst_smooth diff --git a/src/utils/helper.py b/src/utils/helper.py new file mode 100644 index 0000000..cf97033 --- /dev/null +++ b/src/utils/helper.py @@ -0,0 +1,199 @@ +# coding: utf-8 + +""" +utility functions and classes to handle feature extraction and model loading +""" + +import os +import os.path as osp +import torch +from collections import OrderedDict +import numpy as np +from scipy.spatial import ConvexHull # pylint: disable=E0401,E0611 +from typing import Union +import cv2 + +from ..modules.spade_generator import SPADEDecoder +from ..modules.warping_network import WarpingNetwork +from ..modules.motion_extractor import MotionExtractor +from ..modules.appearance_feature_extractor import AppearanceFeatureExtractor +from ..modules.stitching_retargeting_network import StitchingRetargetingNetwork + + +def tensor_to_numpy(data: Union[np.ndarray, torch.Tensor]) -> np.ndarray: + """transform torch.Tensor into numpy.ndarray""" + if isinstance(data, torch.Tensor): + return data.data.cpu().numpy() + return data + +def calc_motion_multiplier( + kp_source: Union[np.ndarray, torch.Tensor], + kp_driving_initial: Union[np.ndarray, torch.Tensor] +) -> float: + """calculate motion_multiplier based on the source image and the first driving frame""" + kp_source_np = tensor_to_numpy(kp_source) + kp_driving_initial_np = tensor_to_numpy(kp_driving_initial) + + source_area = ConvexHull(kp_source_np.squeeze(0)).volume + driving_area = ConvexHull(kp_driving_initial_np.squeeze(0)).volume + motion_multiplier = np.sqrt(source_area) / np.sqrt(driving_area) + # motion_multiplier = np.cbrt(source_area) / np.cbrt(driving_area) + + return motion_multiplier + +def suffix(filename): + """a.jpg -> jpg""" + pos = filename.rfind(".") + if pos == -1: + return "" + return filename[pos + 1:] + + +def prefix(filename): + """a.jpg -> a""" + pos = filename.rfind(".") + if pos == -1: + return filename + return filename[:pos] + + +def basename(filename): + """a/b/c.jpg -> c""" + return prefix(osp.basename(filename)) + + +def remove_suffix(filepath): + """a/b/c.jpg -> a/b/c""" + return osp.join(osp.dirname(filepath), basename(filepath)) + + +def is_image(file_path): + image_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp') + return file_path.lower().endswith(image_extensions) + + +def is_video(file_path): + if file_path.lower().endswith((".mp4", ".mov", ".avi", ".webm")) or osp.isdir(file_path): + return True + return False + + +def is_template(file_path): + if file_path.endswith(".pkl"): + return True + return False + + +def mkdir(d, log=False): + # return self-assined `d`, for one line code + if not osp.exists(d): + os.makedirs(d, exist_ok=True) + if log: + print(f"Make dir: {d}") + return d + + +def squeeze_tensor_to_numpy(tensor): + out = tensor.data.squeeze(0).cpu().numpy() + return out + + +def dct2device(dct: dict, device): + for key in dct: + if isinstance(dct[key], torch.Tensor): + dct[key] = dct[key].to(device) + else: + dct[key] = torch.tensor(dct[key]).to(device) + return dct + + +def concat_feat(kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: + """ + kp_source: (bs, k, 3) + kp_driving: (bs, k, 3) + Return: (bs, 2k*3) + """ + bs_src = kp_source.shape[0] + bs_dri = kp_driving.shape[0] + assert bs_src == bs_dri, 'batch size must be equal' + + feat = torch.cat([kp_source.view(bs_src, -1), kp_driving.view(bs_dri, -1)], dim=1) + return feat + + +def remove_ddp_dumplicate_key(state_dict): + state_dict_new = OrderedDict() + for key in state_dict.keys(): + state_dict_new[key.replace('module.', '')] = state_dict[key] + return state_dict_new + + +def load_model(ckpt_path, model_config, device, model_type): + model_params = model_config['model_params'][f'{model_type}_params'] + + if model_type == 'appearance_feature_extractor': + model = AppearanceFeatureExtractor(**model_params).to(device) + elif model_type == 'motion_extractor': + model = MotionExtractor(**model_params).to(device) + elif model_type == 'warping_module': + model = WarpingNetwork(**model_params).to(device) + elif model_type == 'spade_generator': + model = SPADEDecoder(**model_params).to(device) + elif model_type == 'stitching_retargeting_module': + # Special handling for stitching and retargeting module + config = model_config['model_params']['stitching_retargeting_module_params'] + checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage) + + stitcher = StitchingRetargetingNetwork(**config.get('stitching')) + stitcher.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_shoulder'])) + stitcher = stitcher.to(device) + stitcher.eval() + + retargetor_lip = StitchingRetargetingNetwork(**config.get('lip')) + retargetor_lip.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_mouth'])) + retargetor_lip = retargetor_lip.to(device) + retargetor_lip.eval() + + retargetor_eye = StitchingRetargetingNetwork(**config.get('eye')) + retargetor_eye.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_eye'])) + retargetor_eye = retargetor_eye.to(device) + retargetor_eye.eval() + + return { + 'stitching': stitcher, + 'lip': retargetor_lip, + 'eye': retargetor_eye + } + else: + raise ValueError(f"Unknown model type: {model_type}") + + model.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage)) + model.eval() + return model + + +def load_description(fp): + with open(fp, 'r', encoding='utf-8') as f: + content = f.read() + return content + + +def is_square_video(video_path): + video = cv2.VideoCapture(video_path) + + width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + video.release() + # if width != height: + # gr.Info(f"Uploaded video is not square, force do crop (driving) to be True") + + return width == height + +def clean_state_dict(state_dict): + new_state_dict = OrderedDict() + for k, v in state_dict.items(): + if k[:7] == 'module.': + k = k[7:] # remove `module.` + new_state_dict[k] = v + return new_state_dict diff --git a/src/utils/human_landmark_runner.py b/src/utils/human_landmark_runner.py new file mode 100644 index 0000000..e25e5a4 --- /dev/null +++ b/src/utils/human_landmark_runner.py @@ -0,0 +1,95 @@ +# coding: utf-8 + +import os.path as osp +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) +import torch +import numpy as np +import onnxruntime +from .timer import Timer +from .rprint import rlog +from .crop import crop_image, _transform_pts + + +def make_abs_path(fn): + return osp.join(osp.dirname(osp.realpath(__file__)), fn) + + +def to_ndarray(obj): + if isinstance(obj, torch.Tensor): + return obj.cpu().numpy() + elif isinstance(obj, np.ndarray): + return obj + else: + return np.array(obj) + + +class LandmarkRunner(object): + """landmark runner""" + + def __init__(self, **kwargs): + ckpt_path = kwargs.get('ckpt_path') + onnx_provider = kwargs.get('onnx_provider', 'cuda') # 默认用cuda + device_id = kwargs.get('device_id', 0) + self.dsize = kwargs.get('dsize', 224) + self.timer = Timer() + + if onnx_provider.lower() == 'cuda': + self.session = onnxruntime.InferenceSession( + ckpt_path, providers=[ + ('CUDAExecutionProvider', {'device_id': device_id}) + ] + ) + elif onnx_provider.lower() == 'mps': + self.session = onnxruntime.InferenceSession( + ckpt_path, providers=[ + 'CoreMLExecutionProvider' + ] + ) + else: + opts = onnxruntime.SessionOptions() + opts.intra_op_num_threads = 4 # 默认线程数为 4 + self.session = onnxruntime.InferenceSession( + ckpt_path, providers=['CPUExecutionProvider'], + sess_options=opts + ) + + def _run(self, inp): + out = self.session.run(None, {'input': inp}) + return out + + def run(self, img_rgb: np.ndarray, lmk=None): + if lmk is not None: + crop_dct = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1) + img_crop_rgb = crop_dct['img_crop'] + else: + # NOTE: force resize to 224x224, NOT RECOMMEND! + img_crop_rgb = cv2.resize(img_rgb, (self.dsize, self.dsize)) + scale = max(img_rgb.shape[:2]) / self.dsize + crop_dct = { + 'M_c2o': np.array([ + [scale, 0., 0.], + [0., scale, 0.], + [0., 0., 1.], + ], dtype=np.float32), + } + + inp = (img_crop_rgb.astype(np.float32) / 255.).transpose(2, 0, 1)[None, ...] # HxWx3 (BGR) -> 1x3xHxW (RGB!) + + out_lst = self._run(inp) + out_pts = out_lst[2] + + # 2d landmarks 203 points + lmk = to_ndarray(out_pts[0]).reshape(-1, 2) * self.dsize # scale to 0-224 + lmk = _transform_pts(lmk, M=crop_dct['M_c2o']) + + return lmk + + def warmup(self): + self.timer.tic() + + dummy_image = np.zeros((1, 3, self.dsize, self.dsize), dtype=np.float32) + + _ = self._run(dummy_image) + + elapse = self.timer.toc() + rlog(f'LandmarkRunner warmup time: {elapse:.3f}s') diff --git a/src/utils/io.py b/src/utils/io.py new file mode 100644 index 0000000..28ca742 --- /dev/null +++ b/src/utils/io.py @@ -0,0 +1,117 @@ +# coding: utf-8 + +import os.path as osp +import imageio +import numpy as np +import pickle +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) + +from .helper import mkdir, suffix + + +def load_image_rgb(image_path: str): + if not osp.exists(image_path): + raise FileNotFoundError(f"Image not found: {image_path}") + img = cv2.imread(image_path, cv2.IMREAD_COLOR) + return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + +def load_video(video_info, n_frames=-1): + reader = imageio.get_reader(video_info, "ffmpeg") + + ret = [] + for idx, frame_rgb in enumerate(reader): + if n_frames > 0 and idx >= n_frames: + break + ret.append(frame_rgb) + + reader.close() + return ret + + +def contiguous(obj): + if not obj.flags.c_contiguous: + obj = obj.copy(order="C") + return obj + + +def resize_to_limit(img: np.ndarray, max_dim=1920, division=2): + """ + ajust the size of the image so that the maximum dimension does not exceed max_dim, and the width and the height of the image are multiples of n. + :param img: the image to be processed. + :param max_dim: the maximum dimension constraint. + :param n: the number that needs to be multiples of. + :return: the adjusted image. + """ + h, w = img.shape[:2] + + # ajust the size of the image according to the maximum dimension + if max_dim > 0 and max(h, w) > max_dim: + if h > w: + new_h = max_dim + new_w = int(w * (max_dim / h)) + else: + new_w = max_dim + new_h = int(h * (max_dim / w)) + img = cv2.resize(img, (new_w, new_h)) + + # ensure that the image dimensions are multiples of n + division = max(division, 1) + new_h = img.shape[0] - (img.shape[0] % division) + new_w = img.shape[1] - (img.shape[1] % division) + + if new_h == 0 or new_w == 0: + # when the width or height is less than n, no need to process + return img + + if new_h != img.shape[0] or new_w != img.shape[1]: + img = img[:new_h, :new_w] + + return img + + +def load_img_online(obj, mode="bgr", **kwargs): + max_dim = kwargs.get("max_dim", 1920) + n = kwargs.get("n", 2) + if isinstance(obj, str): + if mode.lower() == "gray": + img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE) + else: + img = cv2.imread(obj, cv2.IMREAD_COLOR) + else: + img = obj + + # Resize image to satisfy constraints + img = resize_to_limit(img, max_dim=max_dim, division=n) + + if mode.lower() == "bgr": + return contiguous(img) + elif mode.lower() == "rgb": + return contiguous(img[..., ::-1]) + else: + raise Exception(f"Unknown mode {mode}") + + +def load(fp): + suffix_ = suffix(fp) + + if suffix_ == "npy": + return np.load(fp) + elif suffix_ == "pkl": + return pickle.load(open(fp, "rb")) + else: + raise Exception(f"Unknown type: {suffix}") + + +def dump(wfp, obj): + wd = osp.split(wfp)[0] + if wd != "" and not osp.exists(wd): + mkdir(wd) + + _suffix = suffix(wfp) + if _suffix == "npy": + np.save(wfp, obj) + elif _suffix == "pkl": + pickle.dump(obj, open(wfp, "wb")) + else: + raise Exception("Unknown type: {}".format(_suffix)) diff --git a/src/utils/resources/clip_embedding_68.pkl b/src/utils/resources/clip_embedding_68.pkl new file mode 100644 index 0000000..d69ca7f Binary files /dev/null and b/src/utils/resources/clip_embedding_68.pkl differ diff --git a/src/utils/resources/clip_embedding_9.pkl b/src/utils/resources/clip_embedding_9.pkl new file mode 100644 index 0000000..2b1e6ae Binary files /dev/null and b/src/utils/resources/clip_embedding_9.pkl differ diff --git a/src/utils/resources/lip_array.pkl b/src/utils/resources/lip_array.pkl new file mode 100644 index 0000000..61fe803 Binary files /dev/null and b/src/utils/resources/lip_array.pkl differ diff --git a/src/utils/resources/mask_template.png b/src/utils/resources/mask_template.png new file mode 100644 index 0000000..bca6ca5 Binary files /dev/null and b/src/utils/resources/mask_template.png differ diff --git a/src/utils/retargeting_utils.py b/src/utils/retargeting_utils.py new file mode 100644 index 0000000..cbe9d1c --- /dev/null +++ b/src/utils/retargeting_utils.py @@ -0,0 +1,24 @@ + +""" +Functions to compute distance ratios between specific pairs of facial landmarks +""" + +import numpy as np + + +def calculate_distance_ratio(lmk: np.ndarray, idx1: int, idx2: int, idx3: int, idx4: int, eps: float = 1e-6) -> np.ndarray: + return (np.linalg.norm(lmk[:, idx1] - lmk[:, idx2], axis=1, keepdims=True) / + (np.linalg.norm(lmk[:, idx3] - lmk[:, idx4], axis=1, keepdims=True) + eps)) + + +def calc_eye_close_ratio(lmk: np.ndarray, target_eye_ratio: np.ndarray = None) -> np.ndarray: + lefteye_close_ratio = calculate_distance_ratio(lmk, 6, 18, 0, 12) + righteye_close_ratio = calculate_distance_ratio(lmk, 30, 42, 24, 36) + if target_eye_ratio is not None: + return np.concatenate([lefteye_close_ratio, righteye_close_ratio, target_eye_ratio], axis=1) + else: + return np.concatenate([lefteye_close_ratio, righteye_close_ratio], axis=1) + + +def calc_lip_close_ratio(lmk: np.ndarray) -> np.ndarray: + return calculate_distance_ratio(lmk, 90, 102, 48, 66) diff --git a/src/utils/rprint.py b/src/utils/rprint.py new file mode 100644 index 0000000..18614e8 --- /dev/null +++ b/src/utils/rprint.py @@ -0,0 +1,16 @@ +# coding: utf-8 + +""" +custom print and log functions +""" + +__all__ = ['rprint', 'rlog'] + +try: + from rich.console import Console + console = Console() + rprint = console.print + rlog = console.log +except: + rprint = print + rlog = print diff --git a/src/utils/timer.py b/src/utils/timer.py new file mode 100644 index 0000000..792af24 --- /dev/null +++ b/src/utils/timer.py @@ -0,0 +1,29 @@ +# coding: utf-8 + +""" +tools to measure elapsed time +""" + +import time + +class Timer(object): + """A simple timer.""" + + def __init__(self): + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.diff = 0. + + def tic(self): + # using time.time instead of time.clock because time time.clock + # does not normalize for multithreading + self.start_time = time.time() + + def toc(self, average=True): + self.diff = time.time() - self.start_time + return self.diff + + def clear(self): + self.start_time = 0. + self.diff = 0. diff --git a/src/utils/video.py b/src/utils/video.py new file mode 100644 index 0000000..9331c90 --- /dev/null +++ b/src/utils/video.py @@ -0,0 +1,218 @@ +# coding: utf-8 + +""" +Functions for processing video + +ATTENTION: you need to install ffmpeg and ffprobe in your env! +""" + +import os.path as osp +import numpy as np +import subprocess +import imageio +import cv2 +from rich.progress import track + +from .rprint import rlog as log +from .rprint import rprint as print +from .helper import prefix + + +def exec_cmd(cmd): + return subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + + +def images2video(images, wfp, **kwargs): + fps = kwargs.get('fps', 30) + video_format = kwargs.get('format', 'mp4') # default is mp4 format + codec = kwargs.get('codec', 'libx264') # default is libx264 encoding + quality = kwargs.get('quality') # video quality + pixelformat = kwargs.get('pixelformat', 'yuv420p') # video pixel format + image_mode = kwargs.get('image_mode', 'rgb') + macro_block_size = kwargs.get('macro_block_size', 2) + ffmpeg_params = ['-crf', str(kwargs.get('crf', 18))] + + writer = imageio.get_writer( + wfp, fps=fps, format=video_format, + codec=codec, quality=quality, ffmpeg_params=ffmpeg_params, pixelformat=pixelformat, macro_block_size=macro_block_size + ) + + n = len(images) + for i in track(range(n), description='Writing', transient=True): + if image_mode.lower() == 'bgr': + writer.append_data(images[i][..., ::-1]) + else: + writer.append_data(images[i]) + + writer.close() + + +def video2gif(video_fp, fps=30, size=256): + if osp.exists(video_fp): + d = osp.split(video_fp)[0] + fn = prefix(osp.basename(video_fp)) + palette_wfp = osp.join(d, 'palette.png') + gif_wfp = osp.join(d, f'{fn}.gif') + # generate the palette + cmd = f'ffmpeg -i "{video_fp}" -vf "fps={fps},scale={size}:-1:flags=lanczos,palettegen" "{palette_wfp}" -y' + exec_cmd(cmd) + # use the palette to generate the gif + cmd = f'ffmpeg -i "{video_fp}" -i "{palette_wfp}" -filter_complex "fps={fps},scale={size}:-1:flags=lanczos[x];[x][1:v]paletteuse" "{gif_wfp}" -y' + exec_cmd(cmd) + return gif_wfp + else: + raise FileNotFoundError(f"video_fp: {video_fp} not exists!") + + +def merge_audio_video(video_fp, audio_fp, wfp): + if osp.exists(video_fp) and osp.exists(audio_fp): + cmd = f'ffmpeg -i "{video_fp}" -i "{audio_fp}" -c:v copy -c:a aac "{wfp}" -y' + exec_cmd(cmd) + print(f'merge {video_fp} and {audio_fp} to {wfp}') + else: + print(f'video_fp: {video_fp} or audio_fp: {audio_fp} not exists!') + + +def blend(img: np.ndarray, mask: np.ndarray, background_color=(255, 255, 255)): + mask_float = mask.astype(np.float32) / 255. + background_color = np.array(background_color).reshape([1, 1, 3]) + bg = np.ones_like(img) * background_color + img = np.clip(mask_float * img + (1 - mask_float) * bg, 0, 255).astype(np.uint8) + return img + + +def concat_frames(driving_image_lst, source_image_lst, I_p_lst): + # TODO: add more concat style, e.g., left-down corner driving + out_lst = [] + h, w, _ = I_p_lst[0].shape + source_image_resized_lst = [cv2.resize(img, (w, h)) for img in source_image_lst] + + for idx, _ in track(enumerate(I_p_lst), total=len(I_p_lst), description='Concatenating result...'): + I_p = I_p_lst[idx] + source_image_resized = source_image_resized_lst[idx] if len(source_image_lst) > 1 else source_image_resized_lst[0] + + if driving_image_lst is None: + out = np.hstack((source_image_resized, I_p)) + else: + driving_image = driving_image_lst[idx] + driving_image_resized = cv2.resize(driving_image, (w, h)) + out = np.hstack((driving_image_resized, source_image_resized, I_p)) + + out_lst.append(out) + return out_lst + + +class VideoWriter: + def __init__(self, **kwargs): + self.fps = kwargs.get('fps', 30) + self.wfp = kwargs.get('wfp', 'video.mp4') + self.video_format = kwargs.get('format', 'mp4') + self.codec = kwargs.get('codec', 'libx264') + self.quality = kwargs.get('quality') + self.pixelformat = kwargs.get('pixelformat', 'yuv420p') + self.image_mode = kwargs.get('image_mode', 'rgb') + self.ffmpeg_params = kwargs.get('ffmpeg_params') + + self.writer = imageio.get_writer( + self.wfp, fps=self.fps, format=self.video_format, + codec=self.codec, quality=self.quality, + ffmpeg_params=self.ffmpeg_params, pixelformat=self.pixelformat + ) + + def write(self, image): + if self.image_mode.lower() == 'bgr': + self.writer.append_data(image[..., ::-1]) + else: + self.writer.append_data(image) + + def close(self): + if self.writer is not None: + self.writer.close() + + +def change_video_fps(input_file, output_file, fps=20, codec='libx264', crf=12): + cmd = f'ffmpeg -i "{input_file}" -c:v {codec} -crf {crf} -r {fps} "{output_file}" -y' + exec_cmd(cmd) + + +def get_fps(filepath, default_fps=25): + try: + fps = cv2.VideoCapture(filepath).get(cv2.CAP_PROP_FPS) + + if fps in (0, None): + fps = default_fps + except Exception as e: + log(e) + fps = default_fps + + return fps + + +def has_audio_stream(video_path: str) -> bool: + """ + Check if the video file contains an audio stream. + + :param video_path: Path to the video file + :return: True if the video contains an audio stream, False otherwise + """ + if osp.isdir(video_path): + return False + + cmd = [ + 'ffprobe', + '-v', 'error', + '-select_streams', 'a', + '-show_entries', 'stream=codec_type', + '-of', 'default=noprint_wrappers=1:nokey=1', + f'"{video_path}"' + ] + + try: + # result = subprocess.run(cmd, capture_output=True, text=True) + result = exec_cmd(' '.join(cmd)) + if result.returncode != 0: + log(f"Error occurred while probing video: {result.stderr}") + return False + + # Check if there is any output from ffprobe command + return bool(result.stdout.strip()) + except Exception as e: + log( + f"Error occurred while probing video: {video_path}, " + "you may need to install ffprobe! (https://ffmpeg.org/download.html) " + "Now set audio to false!", + style="bold red" + ) + return False + + +def add_audio_to_video(silent_video_path: str, audio_video_path: str, output_video_path: str): + cmd = [ + 'ffmpeg', + '-y', + '-i', f'"{silent_video_path}"', + '-i', f'"{audio_video_path}"', + '-map', '0:v', + '-map', '1:a', + '-c:v', 'copy', + '-shortest', + f'"{output_video_path}"' + ] + + try: + exec_cmd(' '.join(cmd)) + log(f"Video with audio generated successfully: {output_video_path}") + except subprocess.CalledProcessError as e: + log(f"Error occurred: {e}") + + +def bb_intersection_over_union(boxA, boxB): + xA = max(boxA[0], boxB[0]) + yA = max(boxA[1], boxB[1]) + xB = min(boxA[2], boxB[2]) + yB = min(boxA[3], boxB[3]) + interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1) + boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1) + boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1) + iou = interArea / float(boxAArea + boxBArea - interArea) + return iou diff --git a/src/utils/viz.py b/src/utils/viz.py new file mode 100644 index 0000000..59b35d2 --- /dev/null +++ b/src/utils/viz.py @@ -0,0 +1,19 @@ +# coding: utf-8 + +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) + + +def viz_lmk(img_, vps, **kwargs): + """可视化点""" + lineType = kwargs.get("lineType", cv2.LINE_8) # cv2.LINE_AA + img_for_viz = img_.copy() + for pt in vps: + cv2.circle( + img_for_viz, + (int(pt[0]), int(pt[1])), + radius=kwargs.get("radius", 1), + color=(0, 255, 0), + thickness=kwargs.get("thickness", 1), + lineType=lineType, + ) + return img_for_viz