Spaces:
Runtime error
Runtime error
| import cv2 | |
| import numpy as np | |
| from ..utils.load_model import load_model | |
| class FaceMesh: | |
| def __init__(self, model_path, device="cuda"): | |
| self.model, self.model_type = load_model(model_path, device=device) | |
| self.input_size = (256, 256) # (w, h) | |
| self.output_names = [ | |
| "Identity", | |
| "Identity_1", | |
| "Identity_2", | |
| ] # Identity is the mesh | |
| def project_landmarks(self, points, roi): | |
| width, height = self.input_size | |
| points /= (width, height, width) | |
| sin, cos = np.sin(roi[4]), np.cos(roi[4]) | |
| matrix = np.array([[cos, sin, 0.0], [-sin, cos, 0.0], [1.0, 1.0, 1.0]]) | |
| points -= (0.5, 0.5, 0.0) | |
| rotated = np.matmul(points * (1, 1, 0), matrix) | |
| points *= (0, 0, 1) | |
| points += rotated | |
| points *= (roi[2], roi[3], roi[2]) | |
| points += (roi[0], roi[1], 0.0) | |
| return points | |
| def __call__(self, image, roi): | |
| """ | |
| image: np.ndarray, RGB, (H, W, C), [0, 255] | |
| roi: np.ndarray, (cx, cy, w, h, rotation), rotation in radian | |
| """ | |
| cx, cy, w, h = roi[:4] | |
| w_half, h_half = w / 2, h / 2 | |
| pts = [ | |
| (cx - w_half, cy - h_half), | |
| (cx + w_half, cy - h_half), | |
| (cx + w_half, cy + h_half), | |
| (cx - w_half, cy + h_half), | |
| ] | |
| rotation = roi[4] | |
| s, c = np.sin(rotation), np.cos(rotation) | |
| t = np.array(pts) - (cx, cy) | |
| r = np.array([[c, s], [-s, c]]) | |
| src_pts = np.matmul(t, r) + (cx, cy) | |
| src_pts = src_pts.astype(np.float32) | |
| dst_pts = np.array( | |
| [ | |
| [0.0, 0.0], | |
| [self.input_size[0], 0.0], | |
| [self.input_size[0], self.input_size[1]], | |
| [0.0, self.input_size[1]], | |
| ] | |
| ).astype(np.float32) | |
| M = cv2.getPerspectiveTransform(src_pts, dst_pts) | |
| roi_image = cv2.warpPerspective( | |
| image, M, self.input_size, flags=cv2.INTER_LINEAR | |
| ) | |
| # cv2.imwrite('test.jpg', cv2.cvtColor(roi_image, cv2.COLOR_RGB2BGR)) | |
| roi_image = roi_image / 255.0 | |
| roi_image = roi_image.astype(np.float32) | |
| roi_image = roi_image[np.newaxis, :, :, :] | |
| outputs = {} | |
| if self.model_type == "onnx": | |
| out_list = self.model.run(None, {"input": roi_image}) | |
| for i, name in enumerate(self.output_names): | |
| outputs[name] = out_list[i] | |
| elif self.model_type == "tensorrt": | |
| self.model.setup({"input": roi_image}) | |
| self.model.infer() | |
| for name in self.output_names: | |
| outputs[name] = self.model.buffer[name][0] | |
| else: | |
| raise ValueError(f"Unsupported model type: {self.model_type}") | |
| points = outputs["Identity"].reshape(1434 // 3, 3) | |
| points = self.project_landmarks(points, roi) | |
| return points | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model", type=str, help="model path") | |
| parser.add_argument("--image", type=str, help="image path") | |
| parser.add_argument("--device", type=str, default="cuda", help="device") | |
| args = parser.parse_args() | |
| face_mesh = FaceMesh(args.model, args.device) | |
| image = cv2.imread(args.image, cv2.IMREAD_COLOR) | |
| image = cv2.resize(image, (256, 256)) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| face_mesh = FaceMesh(args.model, args.device) | |
| roi = np.array([128, 128, 256, 256, np.pi / 2]) | |
| mesh = face_mesh(image, roi) | |
| print(mesh.shape) | |