diff --git a/.gitignore b/.gitignore index b6e4761..6c95386 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +pretrained +character + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/main.py b/main.py index d2415e7..6bfa040 100644 --- a/main.py +++ b/main.py @@ -1,32 +1,3 @@ -# import time -# -# import pyvirtualcam -# import numpy as np -# import cv2 -# import mediapipe as mp -# -# mp_drawing = mp.solutions.drawing_utils -# mp_holistic = mp.solutions.holistic -# -# cap = cv2.VideoCapture(0) -# -# with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic: -# with pyvirtualcam.Camera(width=1280, height=720, fps=30) as cam: -# print(f'Using virtual camera: {cam.device}') -# while cap.isOpened(): -# ret, frame = cap.read() -# frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) -# frame = cv2.resize(frame, dsize=None, fx=720/480, fy=720/480) -# results = holistic.process(frame) -# mp_drawing.draw_landmarks(frame, results.face_landmarks, mp_holistic.FACEMESH_TESSELATION) -# mp_drawing.draw_landmarks(frame, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS) -# mp_drawing.draw_landmarks(frame, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS) -# mp_drawing.draw_landmarks(frame, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS) -# pad_frame = np.zeros((720, 1280, 3), dtype=np.uint8) -# pad_frame[:, (1280-960)//2:(1280-960)//2+960] = frame -# cam.send(pad_frame) -# cam.sleep_until_next_frame() - import time import argparse @@ -50,10 +21,10 @@ from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_IRIS from mediapipe.python.solutions.face_mesh_connections import FACEMESH_TESSELATION from tha2.poser.modes.mode_20 import load_face_morpher, load_face_rotater, load_combiner -from utils import preprocessing_image, postprocessing_image +from utils import preprocessing_image, postprocessing_image, get_distance parser = argparse.ArgumentParser() -parser.add_argument('--debug', action='store_true') +parser.add_argument('--debug', action='store_false') parser.add_argument('--output', type=str, default='window') args = parser.parse_args() @@ -74,11 +45,96 @@ class TalkingAnimeLight(nn.Module): return output_image +def get_pose(facial_landmarks): + iris_r_center = [0, 0, 0] + iris_l_center = [0, 0, 0] + + for idx in list(map(lambda x: x[0], FACEMESH_RIGHT_IRIS)): + pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + iris_r_center[0] += pos_x + iris_r_center[1] += pos_y + iris_r_center[2] += pos_z + + iris_r_center = np.array(iris_r_center) / len(FACEMESH_RIGHT_IRIS) + + for idx in list(map(lambda x: x[0], FACEMESH_LEFT_IRIS)): + pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + iris_l_center[0] += pos_x + iris_l_center[1] += pos_y + iris_l_center[2] += pos_z + + iris_l_center = np.array(iris_l_center) / len(FACEMESH_LEFT_IRIS) + + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], + # [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], (0, 128, 255), 2) + + z_angle = np.arctan2(iris_l_center[1] - iris_r_center[1], iris_l_center[0] - iris_r_center[0]) # * 180 / np.pi + y_angle = np.arctan2(iris_l_center[2] - iris_r_center[2], iris_l_center[0] - iris_r_center[0]) # * 180 / np.pi + + iris_middle_center = np.mean(np.array([iris_l_center, iris_r_center]), axis=0) + mouth_center = [0, 0, 0] + + for idx in list(map(lambda x: x[0], FACEMESH_LIPS)): + pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + mouth_center[0] += pos_x + mouth_center[1] += pos_y + mouth_center[2] += pos_z + + mouth_center = np.array(mouth_center) / len(FACEMESH_LIPS) + mouth_h = facial_landmarks[13].y - facial_landmarks[14].y + mouth_w = facial_landmarks[78].x - (facial_landmarks[409].x + facial_landmarks[375].x) / 2 + mouth_ratio = mouth_h / mouth_w + + x_angle = np.arctan2(iris_middle_center[2] - mouth_center[2], iris_middle_center[1] - mouth_center[1]) # * 180 / np.pi + x_angle = np.arctan2(iris_middle_center[1] - mouth_center[1], iris_middle_center[2] - mouth_center[2]) # * 180 / np.pi + + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_middle_center[0] * w), int(iris_middle_center[1] * h)], + # [int(mouth_center[0] * w), int(mouth_center[1] * h)], (0, 128, 255), 2) + # cv2.putText(frame, f'z angle:{z_angle * 180 / np.pi}', + # [int(iris_middle_center[0] * w), int(iris_middle_center[1] * h) + 30], cv2.FONT_HERSHEY_SIMPLEX, + # 0.5, (255, 0, 0), 1) + + iris_rotation_l_h = get_distance(facial_landmarks[386], facial_landmarks[374]) + iris_rotation_l_w = get_distance(facial_landmarks[382], facial_landmarks[263]) + iris_rotation_r_h = get_distance(facial_landmarks[159], facial_landmarks[145]) + iris_rotation_r_w = get_distance(facial_landmarks[33], facial_landmarks[155]) + + iris_rotation_l_h_temp = np.sqrt( + (iris_l_center[0] - facial_landmarks[386].x) ** 2 + (iris_l_center[1] - facial_landmarks[386].y) ** 2) + iris_rotation_l_w_temp = np.sqrt( + (iris_l_center[0] - facial_landmarks[382].x) ** 2 + (iris_l_center[1] - facial_landmarks[382].y) ** 2) + iris_rotation_r_h_temp = np.sqrt( + (iris_r_center[0] - facial_landmarks[159].x) ** 2 + (iris_r_center[1] - facial_landmarks[159].y) ** 2) + iris_rotation_r_w_temp = np.sqrt( + (iris_r_center[0] - facial_landmarks[33].x) ** 2 + (iris_r_center[1] - facial_landmarks[33].y) ** 2) + + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], + # [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], (0, 128, 128), 3) + # cv2.line(frame, [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], + # [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], (255, 128, 0), 3) + + eye_x_ratio = ((iris_rotation_l_w_temp / iris_rotation_l_w + iris_rotation_r_w_temp / iris_rotation_r_w) - 1) * 3 + eye_y_ratio = ((iris_rotation_l_h_temp / iris_rotation_l_h + iris_rotation_r_h_temp / iris_rotation_r_h) - 1) * 3 + + eye_l_h_temp = 1 - 2 * (facial_landmarks[145].y - facial_landmarks[159].y) / ( + facial_landmarks[155].x - facial_landmarks[33].x) + eye_r_h_temp = 1 - 2 * (facial_landmarks[374].y - facial_landmarks[386].y) / ( + facial_landmarks[263].x - facial_landmarks[382].x) + + return eye_l_h_temp, eye_r_h_temp, mouth_ratio, eye_y_ratio, eye_x_ratio, y_angle, z_angle, x_angle + def main(): model = TalkingAnimeLight().cuda() model = model.eval() model = model.half() img = Image.open("character/sleeping_rat.png") + img = img.resize((256, 256)) input_image = preprocessing_image(img).unsqueeze(0) mp_facemesh = mp.solutions.face_mesh @@ -92,21 +148,14 @@ def main(): if ret is None: raise Exception("Can't find Camera") - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - results = facemesh.process(frame) + input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + results = facemesh.process(input_frame) if results.multi_face_landmarks is None: raise Exception("Fail to initailize parameters") facial_landmarks = results.multi_face_landmarks[0].landmark - # initial parameters - eye_l_h = facial_landmarks[145].y - facial_landmarks[159].y - eye_r_h = facial_landmarks[374].y - facial_landmarks[386].y - - eye_l_initial = facial_landmarks[159].y - eye_r_initial = facial_landmarks[386].y - mouth_eye_vector = torch.empty(1, 27) pose_vector = torch.empty(1, 3) @@ -118,93 +167,142 @@ def main(): mouth_eye_vector = mouth_eye_vector.cuda() pose_vector = pose_vector.cuda() + pose_queue = [] + while cap.isOpened(): ret, frame = cap.read() - results = facemesh.process(frame) + input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + results = facemesh.process(input_frame) if results.multi_face_landmarks is None: continue facial_landmarks = results.multi_face_landmarks[0].landmark + # + # iris_r_center = [0, 0, 0] + # iris_l_center = [0, 0, 0] + # + # for idx in list(map(lambda x: x[0], FACEMESH_RIGHT_IRIS)): + # pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + # iris_r_center[0] += pos_x + # iris_r_center[1] += pos_y + # iris_r_center[2] += pos_z + # + # iris_r_center = np.array(iris_r_center) / len(FACEMESH_RIGHT_IRIS) + # + # for idx in list(map(lambda x: x[0], FACEMESH_LEFT_IRIS)): + # pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + # iris_l_center[0] += pos_x + # iris_l_center[1] += pos_y + # iris_l_center[2] += pos_z + # + # iris_l_center = np.array(iris_l_center) / len(FACEMESH_LEFT_IRIS) + # + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], + # [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], (0, 128, 255), 2) + # + # z_angle = np.arctan2(iris_l_center[1] - iris_r_center[1], + # iris_l_center[0] - iris_r_center[0]) # * 180 / np.pi + # y_angle = np.arctan2(iris_l_center[2] - iris_r_center[2], + # iris_l_center[0] - iris_r_center[0]) # * 180 / np.pi + # + # iris_middle_center = np.mean(np.array([iris_l_center, iris_r_center]), axis=0) + # mouth_center = [0, 0, 0] + # + # for idx in list(map(lambda x: x[0], FACEMESH_LIPS)): + # pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z + # mouth_center[0] += pos_x + # mouth_center[1] += pos_y + # mouth_center[2] += pos_z + # + # mouth_center = np.array(mouth_center) / len(FACEMESH_LIPS) + # mouth_h = facial_landmarks[13].y - facial_landmarks[14].y + # mouth_w = facial_landmarks[78].x - (facial_landmarks[409].x + facial_landmarks[375].x) / 2 + # mouth_ratio = mouth_h / mouth_w + # + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_middle_center[0] * w), int(iris_middle_center[1] * h)], + # [int(mouth_center[0] * w), int(mouth_center[1] * h)], (0, 128, 255), 2) + # cv2.putText(frame, f'z angle:{z_angle * 180 / np.pi}', + # [int(iris_middle_center[0] * w), int(iris_middle_center[1] * h) + 30], + # cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) + # + # iris_rotation_l_h = get_distance(facial_landmarks[386], facial_landmarks[374]) + # iris_rotation_l_w = get_distance(facial_landmarks[382], facial_landmarks[263]) + # iris_rotation_r_h = get_distance(facial_landmarks[159], facial_landmarks[145]) + # iris_rotation_r_w = get_distance(facial_landmarks[33], facial_landmarks[155]) + # + # iris_rotation_l_h_temp = np.sqrt( + # (iris_l_center[0] - facial_landmarks[386].x) ** 2 + (iris_l_center[1] - facial_landmarks[386].y) ** 2) + # iris_rotation_l_w_temp = np.sqrt( + # (iris_l_center[0] - facial_landmarks[382].x) ** 2 + (iris_l_center[1] - facial_landmarks[382].y) ** 2) + # iris_rotation_r_h_temp = np.sqrt( + # (iris_r_center[0] - facial_landmarks[159].x) ** 2 + (iris_r_center[1] - facial_landmarks[159].y) ** 2) + # iris_rotation_r_w_temp = np.sqrt( + # (iris_r_center[0] - facial_landmarks[33].x) ** 2 + (iris_r_center[1] - facial_landmarks[33].y) ** 2) + # + # if args.debug: + # h, w, c = frame.shape + # cv2.line(frame, [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], + # [int(iris_l_center[0] * w), int(iris_l_center[1] * h)], (0, 128, 128), 3) + # cv2.line(frame, [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], + # [int(iris_r_center[0] * w), int(iris_r_center[1] * h)], (255, 128, 0), 3) + # + # eye_x_ratio = (( + # iris_rotation_l_w_temp / iris_rotation_l_w + iris_rotation_r_w_temp / iris_rotation_r_w) - 1) * 3 + # eye_y_ratio = (( + # iris_rotation_l_h_temp / iris_rotation_l_h + iris_rotation_r_h_temp / iris_rotation_r_h) - 1) * 3 + # + # eye_l_h_temp = 0.9 - 2 * (facial_landmarks[145].y - facial_landmarks[159].y) / ( + # facial_landmarks[155].x - facial_landmarks[33].x) + # eye_r_h_temp = 0.9 - 2 * (facial_landmarks[374].y - facial_landmarks[386].y) / ( + # facial_landmarks[263].x - facial_landmarks[382].x) - iris_r_center = [0, 0, 0] - iris_l_center = [0, 0, 0] + pose = get_pose(facial_landmarks) - for idx in list(map(lambda x: x[0], FACEMESH_RIGHT_IRIS)): - pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z - iris_r_center[0] += pos_x - iris_r_center[1] += pos_y - iris_r_center[2] += pos_z + if len(pose_queue) < 3: + pose_queue.append(pose) + pose_queue.append(pose) + pose_queue.append(pose) + else: + pose_queue.pop(0) + pose_queue.append(pose) - iris_r_center = np.array(iris_r_center) / len(FACEMESH_RIGHT_IRIS) + np_pose = np.average(np.array(pose_queue), axis=0, weights=[0.6, 0.3, 0.1]) - for idx in list(map(lambda x: x[0], FACEMESH_LEFT_IRIS)): - pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z - iris_l_center[0] += pos_x - iris_l_center[1] += pos_y - iris_l_center[2] += pos_z - - iris_l_center = np.array(iris_l_center) / len(FACEMESH_LEFT_IRIS) - - z_angle = np.arctan2(iris_l_center[1] - iris_r_center[1], iris_l_center[0] - iris_r_center[0]) #* 180 / np.pi - y_angle = np.arctan2(iris_l_center[2] - iris_r_center[2], iris_l_center[0] - iris_r_center[0]) #* 180 / np.pi - - iris_middle_center = np.mean([iris_l_center, iris_r_center]) - mouth_center = [0, 0, 0] - - for idx in list(map(lambda x: x[0], FACEMESH_LIPS)): - pos_x, pos_y, pos_z = facial_landmarks[idx].x, facial_landmarks[idx].y, facial_landmarks[idx].z - mouth_center[0] += pos_x - mouth_center[1] += pos_y - mouth_center[2] += pos_z - - mouth_center = np.array(mouth_center) / len(FACEMESH_LIPS) - mouth_h = facial_landmarks[13].y - facial_landmarks[14].y - mouth_w = facial_landmarks[78].x - (facial_landmarks[409].x + facial_landmarks[375].x)/2 - mouth_ratio = mouth_h / mouth_w - - eye_l_x_min_boundary = min(facial_landmarks[33].x, facial_landmarks[145].x, facial_landmarks[159].x, facial_landmarks[155].x) - eye_l_x_max_boundary = max(facial_landmarks[33].x, facial_landmarks[145].x, facial_landmarks[159].x, facial_landmarks[155].x) - eye_l_y_min_boundary = min(facial_landmarks[33].y, facial_landmarks[145].y, facial_landmarks[159].y, facial_landmarks[155].y) - eye_l_y_max_boundary = max(facial_landmarks[33].y, facial_landmarks[145].y, facial_landmarks[159].y, facial_landmarks[155].y) - - eye_r_x_min_boundary = min(facial_landmarks[382].x, facial_landmarks[386].x, facial_landmarks[374].x, facial_landmarks[263].x) - eye_r_x_max_boundary = max(facial_landmarks[382].x, facial_landmarks[386].x, facial_landmarks[374].x, facial_landmarks[263].x) - eye_r_y_min_boundary = min(facial_landmarks[382].y, facial_landmarks[386].y, facial_landmarks[374].y, facial_landmarks[263].y) - eye_r_y_max_boundary = max(facial_landmarks[382].y, facial_landmarks[386].y, facial_landmarks[374].y, facial_landmarks[263].y) - - eye_l_x_ratio = (iris_l_center[0] - eye_l_x_min_boundary) / (eye_l_x_max_boundary - eye_l_x_min_boundary) - eye_l_y_ratio = (iris_l_center[1] - eye_l_y_min_boundary) / (eye_l_y_max_boundary - eye_l_y_min_boundary) - - eye_r_x_ratio = (iris_r_center[0] - eye_r_x_min_boundary) / (eye_r_x_max_boundary - eye_r_x_min_boundary) - eye_r_y_ratio = (iris_r_center[1] - eye_r_y_min_boundary) / (eye_r_y_max_boundary - eye_r_y_min_boundary) - - eye_x_ratio = (eye_l_x_ratio + eye_r_x_ratio) / 2 - eye_y_ratio = (eye_l_y_ratio + eye_r_y_ratio) / 2 - - eye_l_h_temp = (facial_landmarks[159].y - eye_l_initial) / eye_l_h - eye_r_h_temp = (facial_landmarks[386].y - eye_r_initial) / eye_r_h - - # eye_l_h_temp = (facial_landmarks[145].x - facial_landmarks[159].x) / (facial_landmarks[155].x - facial_landmarks[33].x) * 5 - # eye_r_h_temp = (facial_landmarks[374].x - facial_landmarks[386].x) / (facial_landmarks[263].x - facial_landmarks[382].x) * 5 + eye_l_h_temp = np_pose[0] + eye_r_h_temp = np_pose[1] + mouth_ratio = np_pose[2] + eye_y_ratio = np_pose[3] + eye_x_ratio = np_pose[4] + y_angle = np_pose[5] + z_angle = np_pose[6] + x_angle = np_pose[7] mouth_eye_vector[0, :] = 0 - mouth_eye_vector[0, 2] = eye_l_h_temp * 1.2 - mouth_eye_vector[0, 3] = eye_r_h_temp * 1.2 + mouth_eye_vector[0, 2] = eye_l_h_temp + mouth_eye_vector[0, 3] = eye_r_h_temp mouth_eye_vector[0, 14] = mouth_ratio * 1.5 - mouth_eye_vector[0, 25] = eye_x_ratio - mouth_eye_vector[0, 26] = eye_y_ratio + mouth_eye_vector[0, 25] = eye_y_ratio + mouth_eye_vector[0, 26] = eye_x_ratio # pose_vector[0, 0] = x_angle - pose_vector[0, 0] = 0 + pose_vector[0, 0] = (1.5 + x_angle) * 2.0 pose_vector[0, 1] = y_angle * 2.0 #temp weight pose_vector[0, 2] = z_angle * 1.5 #temp weight + cv2.putText(frame, f'z angle:{x_angle * 180 / np.pi}, {x_angle}', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) + output_image = model(input_image, mouth_eye_vector, pose_vector) cv2.imshow("frame", cv2.cvtColor(postprocessing_image(output_image.cpu()), cv2.COLOR_RGBA2BGRA)) + if args.debug: + cv2.imshow("camera", frame) cv2.waitKey(1) # if args.output == 'window': @@ -229,5 +327,6 @@ def main(): # x,y,z 27, 28, 29 + if __name__ == '__main__': main() diff --git a/utils.py b/utils.py index 0e03c7e..11d2d0e 100644 --- a/utils.py +++ b/utils.py @@ -1,6 +1,7 @@ import torch import numpy as np + def preprocessing_image(image): np_image = np.array(image) / 255 clipped_image = np.clip(np_image, 0, 1) @@ -27,3 +28,7 @@ def postprocessing_image(image): rgba_image = np.concatenate([rgb_image, alpha_image[..., np.newaxis]], axis=2) rgba_image = rgba_image * 255 return rgba_image.astype(np.uint8) + + +def get_distance(a, b): + return np.sqrt((a.x - b.x) ** 2 + (a.y - b.y) ** 2)