Eye_tracking 1 |

Cargando...
Buscando...
Nada coincide
pygaze.py
Ir a la documentación de este archivo.
1 from omegaconf import OmegaConf
2 from loguru import logger
3 import os
4 import pathlib
5 import torch
6 import cv2
7 import numpy as np
8 from scipy.spatial.transform import Rotation
9 from .utils import get_3d_face_model
10 from .gaze_estimator import GazeEstimator
11 from .common.camera import Camera
12
13 class PyGaze:
14 def __download_model__(self, target_dir):
15 logger.debug('Downloading model to {}...', target_dir)
16 output_dir = pathlib.Path(target_dir).expanduser()
17 output_dir.mkdir(exist_ok=True, parents=True)
18 output_path = os.path.join(output_dir, 'eth-xgaze_resnet18.pth')
19 if not os.path.exists(output_path):
20 logger.debug('Download the pretrained model...')
21 torch.hub.download_url_to_file(
22 'https://github.com/hysts/pytorch_mpiigaze_demo/releases/download/v0.2.2/eth-xgaze_resnet18.pth',
23 output_path)
24 else:
25 logger.debug('The pretrained model {} already exists.', output_path)
26 return output_path
27
28 def __init__(self, device="cpu", model_path = "~/.ptgaze/models"):
29
30 self.config = OmegaConf.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config/eth-xgaze.yaml'))
31 self.config.PACKAGE_ROOT = pathlib.Path(__file__).parent.resolve().as_posix()
32 self.config.device = device
33
34 self.CENTER_X = -0.01821846597470783
35 self.CENTER_Y = -0.17080051119284462
36 self.CENTER_THRESHOLD_X = 0.5
37 self.CENTER_THRESHOLD_Y = 0.3
38
39 # check the model path and download the model
40 self.config.gaze_estimator.checkpoint = os.path.abspath(model_path)
41 if os.path.isfile(self.config.gaze_estimator.checkpoint):
42 logger.warning("{} is a file path but a directory is required. Removing the filename...", self.config.gaze_estimator.checkpoint)
43 self.config.gaze_estimator.checkpoint = os.path.dirname(self.config.gaze_estimator.checkpoint)
44 self.config.gaze_estimator.checkpoint = self.__download_model__(self.config.gaze_estimator.checkpoint)
45
46 # initialize
47 self.gaze_estimator = GazeEstimator(self.config)
48
49 def predict(self, img):
50 results = []
51 if img is None:
52 logger.warning("Invalid image.")
53 return results
54
55 undistorted = cv2.undistort(img, self.gaze_estimator.camera.camera_matrix, self.gaze_estimator.camera.dist_coefficients)
56 faces = self.gaze_estimator.detect_faces(undistorted)
57 for face in faces:
58 self.gaze_estimator.estimate_gaze(undistorted, face)
59 return faces
60
61 def look_at_camera(self, face):
62 return face.gaze_vector[0] > self.CENTER_X - self.CENTER_THRESHOLD_X and face.gaze_vector[0] < self.CENTER_X + self.CENTER_THRESHOLD_X and face.gaze_vector[1] > self.CENTER_Y - self.CENTER_THRESHOLD_Y and face.gaze_vector[1] < self.CENTER_Y + self.CENTER_THRESHOLD_Y
63
64 class PyGazeRenderer:
65 def __init__(self):
66 self.head_pose_axis_length = 0.05
67 self.gaze_visualization_length = 0.05
68 self.AXIS_COLORS = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]
69 config = OmegaConf.load(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config/eth-xgaze.yaml'))
70 config.PACKAGE_ROOT = pathlib.Path(__file__).parent.resolve().as_posix()
71 self.face_3d_model = get_3d_face_model(config)
72 self.camera = Camera(config.gaze_estimator.camera_params)
73
74 def render(self, img, face, draw_face_bbox=True, draw_face_landmarks=True, draw_3dface_model=True,draw_head_pose=True, draw_gaze_vector=True,
75 color = (0, 255, 0)):
76 size = 1
77
78 if draw_face_bbox:
79 bbox = np.round(face.bbox).astype(int).tolist()
80 cv2.rectangle(img, tuple(bbox[0]), tuple(bbox[1]), color, size)
81
82 if draw_face_landmarks:
83 for pt in face.landmarks:
84 pt = tuple(np.round(pt).astype(np.int).tolist())
85 cv2.circle(img, pt, size, color, cv2.FILLED)
86
87 if draw_3dface_model:
88 points2d = self.camera.project_points(face.model3d)
89 for pt in points2d:
90 pt = tuple(np.round(pt).astype(np.int).tolist())
91 cv2.circle(img, pt, size, color, cv2.FILLED)
92
93 if draw_head_pose:
94 axes3d = np.eye(3, dtype=np.float) @ Rotation.from_euler('XYZ', [0, np.pi, 0]).as_matrix()
95 axes3d = axes3d * self.head_pose_axis_length
96 axes2d = self.camera.project_points(axes3d, face.head_pose_rot.as_rotvec(), face.head_position)
97 center = face.landmarks[self.face_3d_model.NOSE_INDEX]
98 center = tuple(np.round(center).astype(np.int).tolist())
99 for pt, color in zip(axes2d, self.AXIS_COLORS):
100 pt = tuple(np.round(pt).astype(np.int).tolist())
101 cv2.line(img, center, pt, color, 2, cv2.LINE_AA)
102
103 if draw_gaze_vector:
104 start = face.center
105 end = face.center + self.gaze_visualization_length * face.gaze_vector
106 points3d = np.vstack([start, end])
107 points2d = self.camera.project_points(points3d)
108 pt0 = tuple(np.round(points2d[0]).astype(int).tolist())
109 pt1 = tuple(np.round(points2d[1]).astype(int).tolist())
110 cv2.line(img, pt0, pt1, color, 1, cv2.LINE_AA)
111
pygaze.gaze_estimator.GazeEstimator
Definition gaze_estimator.py:17
pygaze.pygaze.PyGaze
Definition pygaze.py:13
pygaze.pygaze.PyGaze.__init__
__init__(self, device="cpu", model_path="~/.ptgaze/models")
Definition pygaze.py:28
pygaze.pygaze.PyGaze.predict
predict(self, img)
Definition pygaze.py:49
pygaze.pygaze.PyGaze.look_at_camera
look_at_camera(self, face)
Definition pygaze.py:61
pygaze.pygaze.PyGaze.config
config
Definition pygaze.py:30
pygaze.pygaze.PyGaze.CENTER_X
float CENTER_X
Definition pygaze.py:34
pygaze.pygaze.PyGaze.CENTER_THRESHOLD_X
float CENTER_THRESHOLD_X
Definition pygaze.py:36
pygaze.pygaze.PyGaze.CENTER_THRESHOLD_Y
float CENTER_THRESHOLD_Y
Definition pygaze.py:37
pygaze.pygaze.PyGaze.__download_model__
__download_model__(self, target_dir)
Definition pygaze.py:14
pygaze.pygaze.PyGaze.CENTER_Y
float CENTER_Y
Definition pygaze.py:35
pygaze.pygaze.PyGaze.gaze_estimator
gaze_estimator
Definition pygaze.py:47
pygaze.pygaze.PyGazeRenderer
Definition pygaze.py:64
pygaze.pygaze.PyGazeRenderer.camera
camera
Definition pygaze.py:72
pygaze.pygaze.PyGazeRenderer.gaze_visualization_length
float gaze_visualization_length
Definition pygaze.py:67
pygaze.pygaze.PyGazeRenderer.head_pose_axis_length
float head_pose_axis_length
Definition pygaze.py:66
pygaze.pygaze.PyGazeRenderer.__init__
__init__(self)
Definition pygaze.py:65
pygaze.pygaze.PyGazeRenderer.AXIS_COLORS
list AXIS_COLORS
Definition pygaze.py:68
pygaze.pygaze.PyGazeRenderer.render
render(self, img, face, draw_face_bbox=True, draw_face_landmarks=True, draw_3dface_model=True, draw_head_pose=True, draw_gaze_vector=True, color=(0, 255, 0))
Definition pygaze.py:75
pygaze.pygaze.PyGazeRenderer.face_3d_model
face_3d_model
Definition pygaze.py:71