Eye_tracking 1 |

Cargando...
Buscando...
Nada coincide
gaze_estimator.py
Ir a la documentación de este archivo.
1 from loguru import logger
2 from typing import List
3
4 import numpy as np
5 import torch
6 from omegaconf import DictConfig
7
8 from .common.camera import Camera
9 from .common.face import Face
10 from .common.face_parts import FacePartsName
11 from .head_pose_estimation.head_pose_normalizer import HeadPoseNormalizer
12 from .head_pose_estimation.face_landmark_estimator import LandmarkEstimator
13 from .models import create_model
14 from .transforms import create_transform
15 from .common.face_model_mediapipe import FaceModelMediaPipe
16
17 class GazeEstimator:
18 EYE_KEYS = [FacePartsName.REYE, FacePartsName.LEYE]
19
20 def __init__(self, config: DictConfig):
21 self._config = config
22
23 self._face_model3d = FaceModelMediaPipe()
24
25 self.camera = Camera(config.gaze_estimator.camera_params)
26 self._normalized_camera = Camera(config.gaze_estimator.normalized_camera_params)
27
28 self._landmark_estimator = LandmarkEstimator(config)
29 self._head_pose_normalizer = HeadPoseNormalizer(
30 self.camera, self._normalized_camera,
31 self._config.gaze_estimator.normalized_camera_distance)
32 self._gaze_estimation_model = self._load_model()
33 self._transform = create_transform(config)
34
35 def _load_model(self) -> torch.nn.Module:
36 model = create_model(self._config)
37 checkpoint = torch.load(self._config.gaze_estimator.checkpoint,
38 map_location='cpu')
39 model.load_state_dict(checkpoint['model'])
40 model.to(torch.device(self._config.device))
41 model.eval()
42 return model
43
44 def detect_faces(self, image: np.ndarray) -> List[Face]:
45 return self._landmark_estimator.detect_faces(image)
46
47 def estimate_gaze(self, image: np.ndarray, face: Face) -> None:
48 self._face_model3d.estimate_head_pose(face, self.camera)
49 self._face_model3d.compute_3d_pose(face)
50 self._face_model3d.compute_face_eye_centers(face, self._config.mode)
51 self._head_pose_normalizer.normalize(image, face)
52 self._run_ethxgaze_model(face)
53
54 @torch.no_grad()
55 def _run_ethxgaze_model(self, face: Face) -> None:
56 image = self._transform(face.normalized_image).unsqueeze(0)
57
58 device = torch.device(self._config.device)
59 image = image.to(device)
60 prediction = self._gaze_estimation_model(image)
61 prediction = prediction.cpu().numpy()
62
63 face.normalized_gaze_angles = prediction[0]
64 face.angle_to_vector()
65 face.denormalize_gaze_vector()
pygaze.gaze_estimator.GazeEstimator
Definition gaze_estimator.py:17
pygaze.gaze_estimator.GazeEstimator._face_model3d
_face_model3d
Definition gaze_estimator.py:23
pygaze.gaze_estimator.GazeEstimator._gaze_estimation_model
torch.nn.Module _gaze_estimation_model
Definition gaze_estimator.py:32
pygaze.gaze_estimator.GazeEstimator._head_pose_normalizer
_head_pose_normalizer
Definition gaze_estimator.py:29
pygaze.gaze_estimator.GazeEstimator._transform
_transform
Definition gaze_estimator.py:33
pygaze.gaze_estimator.GazeEstimator._load_model
torch.nn.Module _load_model(self)
Definition gaze_estimator.py:35
pygaze.gaze_estimator.GazeEstimator.estimate_gaze
None estimate_gaze(self, np.ndarray image, Face face)
Definition gaze_estimator.py:47
pygaze.gaze_estimator.GazeEstimator._landmark_estimator
_landmark_estimator
Definition gaze_estimator.py:28
pygaze.gaze_estimator.GazeEstimator._config
_config
Definition gaze_estimator.py:21
pygaze.gaze_estimator.GazeEstimator.__init__
__init__(self, DictConfig config)
Definition gaze_estimator.py:20
pygaze.gaze_estimator.GazeEstimator.camera
camera
Definition gaze_estimator.py:25
pygaze.gaze_estimator.GazeEstimator.detect_faces
List[Face] detect_faces(self, np.ndarray image)
Definition gaze_estimator.py:44
pygaze.gaze_estimator.GazeEstimator._run_ethxgaze_model
None _run_ethxgaze_model(self, Face face)
Definition gaze_estimator.py:55
pygaze.gaze_estimator.GazeEstimator._normalized_camera
_normalized_camera
Definition gaze_estimator.py:26