feat: stage 2 head pose, eye behaviour, MAR/yawn, tighter focus; add torch deps

models/eye_behaviour/eye_attention_model.py

@@ -1 +1,48 @@
-# stub
+# MobileNetV2 eye attention classifier (attentive vs inattentive)
+
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+
+class EyeAttentionModel(nn.Module):
+    def __init__(
+        self,
+        pretrained: bool = True,
+        dropout1: float = 0.3,
+        dropout2: float = 0.2,
+    ):
+        super().__init__()
+
+        weights = models.MobileNet_V2_Weights.DEFAULT if pretrained else None
+        backbone = models.mobilenet_v2(weights=weights)
+
+        self.features = backbone.features
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout1),
+            nn.Linear(1280, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout2),
+            nn.Linear(256, 2),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.pool(x).flatten(1)
+        return self.classifier(x)
+
+    def predict_score(self, x: torch.Tensor) -> torch.Tensor:
+        logits = self.forward(x)
+        probs = torch.softmax(logits, dim=1)
+        return probs[:, 1]
+
+    def freeze_backbone(self):
+        for param in self.features.parameters():
+            param.requires_grad = False
+
+    def unfreeze_last_blocks(self, n: int = 4):
+        total_blocks = len(self.features)
+        for i in range(max(0, total_blocks - n), total_blocks):
+            for param in self.features[i].parameters():
+                param.requires_grad = True

models/eye_behaviour/eye_crop.py

@@ -1 +1,70 @@
-# stub
+# Eye region extraction from Face Mesh landmarks
+
+import cv2
+import numpy as np
+
+LEFT_EYE_CONTOUR = [33, 7, 163, 144, 145, 153, 154, 155, 133, 173, 157, 158, 159, 160, 161, 246]
+RIGHT_EYE_CONTOUR = [362, 382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386, 385, 384, 398]
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+CROP_SIZE = 96
+
+
+def _bbox_from_landmarks(
+    landmarks: np.ndarray,
+    indices: list[int],
+    frame_w: int,
+    frame_h: int,
+    expand: float = 0.4,
+) -> tuple[int, int, int, int]:
+    pts = landmarks[indices, :2]
+    px = pts[:, 0] * frame_w
+    py = pts[:, 1] * frame_h
+
+    x_min, x_max = px.min(), px.max()
+    y_min, y_max = py.min(), py.max()
+    w = x_max - x_min
+    h = y_max - y_min
+    cx = (x_min + x_max) / 2
+    cy = (y_min + y_max) / 2
+
+    size = max(w, h) * (1 + expand)
+    half = size / 2
+
+    x1 = int(max(cx - half, 0))
+    y1 = int(max(cy - half, 0))
+    x2 = int(min(cx + half, frame_w))
+    y2 = int(min(cy + half, frame_h))
+
+    return x1, y1, x2, y2
+
+
+def extract_eye_crops(
+    frame: np.ndarray,
+    landmarks: np.ndarray,
+    expand: float = 0.4,
+    crop_size: int = CROP_SIZE,
+) -> tuple[np.ndarray, np.ndarray, tuple, tuple]:
+    h, w = frame.shape[:2]
+
+    left_bbox = _bbox_from_landmarks(landmarks, LEFT_EYE_CONTOUR, w, h, expand)
+    right_bbox = _bbox_from_landmarks(landmarks, RIGHT_EYE_CONTOUR, w, h, expand)
+
+    left_crop = frame[left_bbox[1] : left_bbox[3], left_bbox[0] : left_bbox[2]]
+    right_crop = frame[right_bbox[1] : right_bbox[3], right_bbox[0] : right_bbox[2]]
+
+    left_crop = cv2.resize(left_crop, (crop_size, crop_size), interpolation=cv2.INTER_AREA)
+    right_crop = cv2.resize(right_crop, (crop_size, crop_size), interpolation=cv2.INTER_AREA)
+
+    return left_crop, right_crop, left_bbox, right_bbox
+
+
+def crop_to_tensor(crop_bgr: np.ndarray):
+    import torch
+
+    rgb = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+    for c in range(3):
+        rgb[:, :, c] = (rgb[:, :, c] - IMAGENET_MEAN[c]) / IMAGENET_STD[c]
+    return torch.from_numpy(rgb.transpose(2, 0, 1))

models/eye_behaviour/eye_scorer.py

@@ -1 +1,167 @@
-# stub
+# EAR + gaze from landmarks -> S_eye (no model)
+
+import math
+
+import numpy as np
+
+_LEFT_EYE_EAR = [33, 160, 158, 133, 153, 145]
+_RIGHT_EYE_EAR = [362, 385, 387, 263, 373, 380]
+
+_LEFT_IRIS_CENTER = 468
+_RIGHT_IRIS_CENTER = 473
+
+_LEFT_EYE_INNER = 133
+_LEFT_EYE_OUTER = 33
+_RIGHT_EYE_INNER = 362
+_RIGHT_EYE_OUTER = 263
+
+_LEFT_EYE_TOP = 159
+_LEFT_EYE_BOTTOM = 145
+_RIGHT_EYE_TOP = 386
+_RIGHT_EYE_BOTTOM = 374
+
+# Mouth (MAR) — inner lip landmarks
+_MOUTH_TOP = 13
+_MOUTH_BOTTOM = 14
+_MOUTH_LEFT = 78
+_MOUTH_RIGHT = 308
+_MOUTH_UPPER_1 = 82
+_MOUTH_UPPER_2 = 312
+_MOUTH_LOWER_1 = 87
+_MOUTH_LOWER_2 = 317
+
+MAR_YAWN_THRESHOLD = 0.55  # MAR above this = mouth open (e.g. yawning / sleepy)
+
+
+def _distance(p1: np.ndarray, p2: np.ndarray) -> float:
+    return float(np.linalg.norm(p1 - p2))
+
+
+def compute_ear(landmarks: np.ndarray, eye_indices: list[int]) -> float:
+    p1 = landmarks[eye_indices[0], :2]
+    p2 = landmarks[eye_indices[1], :2]
+    p3 = landmarks[eye_indices[2], :2]
+    p4 = landmarks[eye_indices[3], :2]
+    p5 = landmarks[eye_indices[4], :2]
+    p6 = landmarks[eye_indices[5], :2]
+
+    vertical1 = _distance(p2, p6)
+    vertical2 = _distance(p3, p5)
+    horizontal = _distance(p1, p4)
+
+    if horizontal < 1e-6:
+        return 0.0
+
+    return (vertical1 + vertical2) / (2.0 * horizontal)
+
+
+def compute_avg_ear(landmarks: np.ndarray) -> float:
+    left_ear = compute_ear(landmarks, _LEFT_EYE_EAR)
+    right_ear = compute_ear(landmarks, _RIGHT_EYE_EAR)
+    return (left_ear + right_ear) / 2.0
+
+
+def compute_gaze_ratio(landmarks: np.ndarray) -> tuple[float, float]:
+    left_iris = landmarks[_LEFT_IRIS_CENTER, :2]
+    left_inner = landmarks[_LEFT_EYE_INNER, :2]
+    left_outer = landmarks[_LEFT_EYE_OUTER, :2]
+    left_top = landmarks[_LEFT_EYE_TOP, :2]
+    left_bottom = landmarks[_LEFT_EYE_BOTTOM, :2]
+
+    right_iris = landmarks[_RIGHT_IRIS_CENTER, :2]
+    right_inner = landmarks[_RIGHT_EYE_INNER, :2]
+    right_outer = landmarks[_RIGHT_EYE_OUTER, :2]
+    right_top = landmarks[_RIGHT_EYE_TOP, :2]
+    right_bottom = landmarks[_RIGHT_EYE_BOTTOM, :2]
+
+    left_h_total = _distance(left_inner, left_outer)
+    right_h_total = _distance(right_inner, right_outer)
+
+    if left_h_total < 1e-6 or right_h_total < 1e-6:
+        return 0.5, 0.5
+
+    left_h_ratio = _distance(left_outer, left_iris) / left_h_total
+    right_h_ratio = _distance(right_outer, right_iris) / right_h_total
+    h_ratio = (left_h_ratio + right_h_ratio) / 2.0
+
+    left_v_total = _distance(left_top, left_bottom)
+    right_v_total = _distance(right_top, right_bottom)
+
+    if left_v_total < 1e-6 or right_v_total < 1e-6:
+        return h_ratio, 0.5
+
+    left_v_ratio = _distance(left_top, left_iris) / left_v_total
+    right_v_ratio = _distance(right_top, right_iris) / right_v_total
+    v_ratio = (left_v_ratio + right_v_ratio) / 2.0
+
+    return float(np.clip(h_ratio, 0, 1)), float(np.clip(v_ratio, 0, 1))
+
+
+def compute_mar(landmarks: np.ndarray) -> float:
+    # Mouth aspect ratio: high = mouth open (yawning / sleepy)
+    top = landmarks[_MOUTH_TOP, :2]
+    bottom = landmarks[_MOUTH_BOTTOM, :2]
+    left = landmarks[_MOUTH_LEFT, :2]
+    right = landmarks[_MOUTH_RIGHT, :2]
+    upper1 = landmarks[_MOUTH_UPPER_1, :2]
+    lower1 = landmarks[_MOUTH_LOWER_1, :2]
+    upper2 = landmarks[_MOUTH_UPPER_2, :2]
+    lower2 = landmarks[_MOUTH_LOWER_2, :2]
+
+    horizontal = _distance(left, right)
+    if horizontal < 1e-6:
+        return 0.0
+    v1 = _distance(upper1, lower1)
+    v2 = _distance(top, bottom)
+    v3 = _distance(upper2, lower2)
+    return (v1 + v2 + v3) / (2.0 * horizontal)
+
+
+class EyeBehaviourScorer:
+    def __init__(
+        self,
+        ear_open: float = 0.30,
+        ear_closed: float = 0.16,
+        gaze_max_offset: float = 0.28,
+    ):
+        self.ear_open = ear_open
+        self.ear_closed = ear_closed
+        self.gaze_max_offset = gaze_max_offset
+
+    def _ear_score(self, ear: float) -> float:
+        if ear >= self.ear_open:
+            return 1.0
+        if ear <= self.ear_closed:
+            return 0.0
+        return (ear - self.ear_closed) / (self.ear_open - self.ear_closed)
+
+    def _gaze_score(self, h_ratio: float, v_ratio: float) -> float:
+        h_offset = abs(h_ratio - 0.5)
+        v_offset = abs(v_ratio - 0.5)
+        offset = math.sqrt(h_offset**2 + v_offset**2)
+        t = min(offset / self.gaze_max_offset, 1.0)
+        return 0.5 * (1.0 + math.cos(math.pi * t))
+
+    def score(self, landmarks: np.ndarray) -> float:
+        ear = compute_avg_ear(landmarks)
+        ear_s = self._ear_score(ear)
+        if ear_s < 0.3:
+            return ear_s
+        h_ratio, v_ratio = compute_gaze_ratio(landmarks)
+        gaze_s = self._gaze_score(h_ratio, v_ratio)
+        return ear_s * gaze_s
+
+    def detailed_score(self, landmarks: np.ndarray) -> dict:
+        ear = compute_avg_ear(landmarks)
+        ear_s = self._ear_score(ear)
+        h_ratio, v_ratio = compute_gaze_ratio(landmarks)
+        gaze_s = self._gaze_score(h_ratio, v_ratio)
+        s_eye = ear_s if ear_s < 0.3 else ear_s * gaze_s
+        return {
+            "ear": round(ear, 4),
+            "ear_score": round(ear_s, 4),
+            "h_gaze": round(h_ratio, 4),
+            "v_gaze": round(v_ratio, 4),
+            "gaze_score": round(gaze_s, 4),
+            "s_eye": round(s_eye, 4),
+        }

models/face_orientation/head_pose.py

@@ -1 +1,114 @@
-# stub
+# Head pose from 6 Face Mesh landmarks (solvePnP) -> yaw/pitch/roll, S_face
+
+import math
+
+import cv2
+import numpy as np
+
+_LANDMARK_INDICES = [1, 152, 33, 263, 61, 291]
+
+_MODEL_POINTS = np.array(
+    [
+        [0.0, 0.0, 0.0],
+        [0.0, -330.0, -65.0],
+        [-225.0, 170.0, -135.0],
+        [225.0, 170.0, -135.0],
+        [-150.0, -150.0, -125.0],
+        [150.0, -150.0, -125.0],
+    ],
+    dtype=np.float64,
+)
+
+
+class HeadPoseEstimator:
+    def __init__(self, max_angle: float = 30.0, roll_weight: float = 0.5):
+        self.max_angle = max_angle
+        self.roll_weight = roll_weight
+        self._camera_matrix = None
+        self._frame_size = None
+        self._dist_coeffs = np.zeros((4, 1), dtype=np.float64)
+
+    def _get_camera_matrix(self, frame_w: int, frame_h: int) -> np.ndarray:
+        if self._camera_matrix is not None and self._frame_size == (frame_w, frame_h):
+            return self._camera_matrix
+        focal_length = float(frame_w)
+        cx, cy = frame_w / 2.0, frame_h / 2.0
+        self._camera_matrix = np.array(
+            [[focal_length, 0, cx], [0, focal_length, cy], [0, 0, 1]],
+            dtype=np.float64,
+        )
+        self._frame_size = (frame_w, frame_h)
+        return self._camera_matrix
+
+    def _solve(self, landmarks: np.ndarray, frame_w: int, frame_h: int):
+        image_points = np.array(
+            [
+                [landmarks[i, 0] * frame_w, landmarks[i, 1] * frame_h]
+                for i in _LANDMARK_INDICES
+            ],
+            dtype=np.float64,
+        )
+        camera_matrix = self._get_camera_matrix(frame_w, frame_h)
+        success, rvec, tvec = cv2.solvePnP(
+            _MODEL_POINTS,
+            image_points,
+            camera_matrix,
+            self._dist_coeffs,
+            flags=cv2.SOLVEPNP_ITERATIVE,
+        )
+        return success, rvec, tvec, image_points
+
+    def estimate(
+        self, landmarks: np.ndarray, frame_w: int, frame_h: int
+    ) -> tuple[float, float, float] | None:
+        success, rvec, tvec, _ = self._solve(landmarks, frame_w, frame_h)
+        if not success:
+            return None
+
+        rmat, _ = cv2.Rodrigues(rvec)
+        nose_dir = rmat @ np.array([0.0, 0.0, 1.0])
+        face_up = rmat @ np.array([0.0, 1.0, 0.0])
+
+        yaw = math.degrees(math.atan2(nose_dir[0], -nose_dir[2]))
+        pitch = math.degrees(math.asin(np.clip(-nose_dir[1], -1.0, 1.0)))
+        roll = math.degrees(math.atan2(face_up[0], -face_up[1]))
+
+        return (yaw, pitch, roll)
+
+    def score(self, landmarks: np.ndarray, frame_w: int, frame_h: int) -> float:
+        angles = self.estimate(landmarks, frame_w, frame_h)
+        if angles is None:
+            return 0.0
+
+        yaw, pitch, roll = angles
+        deviation = math.sqrt(yaw**2 + pitch**2 + (self.roll_weight * roll) ** 2)
+        t = min(deviation / self.max_angle, 1.0)
+        return 0.5 * (1.0 + math.cos(math.pi * t))
+
+    def draw_axes(
+        self,
+        frame: np.ndarray,
+        landmarks: np.ndarray,
+        axis_length: float = 50.0,
+    ) -> np.ndarray:
+        h, w = frame.shape[:2]
+        success, rvec, tvec, image_points = self._solve(landmarks, w, h)
+        if not success:
+            return frame
+
+        camera_matrix = self._get_camera_matrix(w, h)
+        nose = tuple(image_points[0].astype(int))
+
+        axes_3d = np.float64(
+            [[axis_length, 0, 0], [0, axis_length, 0], [0, 0, axis_length]]
+        )
+        projected, _ = cv2.projectPoints(
+            axes_3d, rvec, tvec, camera_matrix, self._dist_coeffs
+        )
+
+        colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]
+        for i, color in enumerate(colors):
+            pt = tuple(projected[i].ravel().astype(int))
+            cv2.line(frame, nose, pt, color, 2)
+
+        return frame

requirements.txt

@@ -1,4 +1,6 @@
-# Stage 1: face mesh + test UI (no torch)
+# Face mesh + head pose + eye behaviour (Stage 2); eye CNN needs torch
 mediapipe>=0.10.14
 opencv-python>=4.8.0
 numpy>=1.24.0
+torch>=2.0.0
+torchvision>=0.15.0

ui/live_demo.py

@@ -119,12 +119,16 @@ def draw_eyes_and_irises(frame, landmarks, w, h):
 
 
 def main():
-    parser = argparse.ArgumentParser(description="FocusGuard — Face mesh (Stage 1)")
+    parser = argparse.ArgumentParser(description="FocusGuard — Face mesh + focus (Stage 2)")
     parser.add_argument("--camera", type=int, default=0, help="Camera index")
+    parser.add_argument("--max-angle", type=float, default=22.0, help="Max head angle for S_face (deg), smaller = tighter")
+    parser.add_argument("--alpha", type=float, default=0.4, help="S_face weight")
+    parser.add_argument("--beta", type=float, default=0.6, help="S_eye weight")
+    parser.add_argument("--threshold", type=float, default=0.55, help="Score >= this = FOCUSED (higher = stricter)")
     args = parser.parse_args()
 
-    print("[DEMO] Face mesh only (Stage 1)")
-    pipeline = FaceMeshPipeline()
+    print("[DEMO] Face mesh + head pose + eye behaviour (Stage 2)")
+    pipeline = FaceMeshPipeline(max_angle=args.max_angle, alpha=args.alpha, beta=args.beta, threshold=args.threshold)
 
     cap = cv2.VideoCapture(args.camera)
     if not cap.isOpened():
@@ -156,10 +160,21 @@ def main():
                 elif mesh_mode == MESH_CONTOURS:
                     draw_contours(frame, lm, w, h)
                 draw_eyes_and_irises(frame, lm, w, h)
-
-            cv2.rectangle(frame, (0, 0), (w, 28), (0, 0, 0), -1)
-            cv2.putText(frame, f"{_MESH_NAMES[mesh_mode]}  FPS: {fps:.0f}", (10, 20), FONT, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(frame, "q:quit  m:mesh", (w - 140, 20), FONT, 0.4, (180, 180, 180), 1, cv2.LINE_AA)
+                pipeline.head_pose.draw_axes(frame, lm)
+
+            # Status bar: FOCUSED / NOT FOCUSED; YAWN when mouth open (sleepy)
+            status = "FOCUSED" if result["is_focused"] else "NOT FOCUSED"
+            status_color = GREEN if result["is_focused"] else RED
+            cv2.rectangle(frame, (0, 0), (w, 55), (0, 0, 0), -1)
+            cv2.putText(frame, status, (10, 28), FONT, 0.8, status_color, 2, cv2.LINE_AA)
+            mar_str = f" MAR:{result['mar']:.2f}" if result.get("mar") is not None else ""
+            cv2.putText(frame, f"S_face:{result['s_face']:.2f} S_eye:{result['s_eye']:.2f}{mar_str} score:{result['raw_score']:.2f}", (10, 48), FONT, 0.45, WHITE, 1, cv2.LINE_AA)
+            if result.get("is_yawning"):
+                cv2.putText(frame, "YAWN", (10, 75), FONT, 0.7, ORANGE, 2, cv2.LINE_AA)
+            if result["yaw"] is not None:
+                cv2.putText(frame, f"yaw:{result['yaw']:+.0f} pitch:{result['pitch']:+.0f} roll:{result['roll']:+.0f}", (w - 280, 48), FONT, 0.4, (180, 180, 180), 1, cv2.LINE_AA)
+            cv2.putText(frame, f"{_MESH_NAMES[mesh_mode]}  FPS: {fps:.0f}", (w - 200, 28), FONT, 0.45, WHITE, 1, cv2.LINE_AA)
+            cv2.putText(frame, "q:quit  m:mesh", (w - 140, 48), FONT, 0.4, (180, 180, 180), 1, cv2.LINE_AA)
 
             cv2.imshow("FocusGuard", frame)
 

ui/pipeline.py

@@ -1,4 +1,4 @@
-# Stage 1: face mesh only (no head pose / eye model / fusion)
+# Stage 2: face mesh + head pose (S_face) + eye behaviour (S_eye) -> focus
 
 import os
 import sys
@@ -10,17 +10,59 @@ if _PROJECT_ROOT not in sys.path:
     sys.path.insert(0, _PROJECT_ROOT)
 
 from models.face_mesh.face_mesh import FaceMeshDetector
+from models.face_orientation.head_pose import HeadPoseEstimator
+from models.eye_behaviour.eye_scorer import EyeBehaviourScorer, compute_mar, MAR_YAWN_THRESHOLD
 
 
 class FaceMeshPipeline:
-    # frame -> face mesh -> 478 landmarks
+    # frame -> face mesh -> S_face + S_eye -> focused / not focused
 
-    def __init__(self):
+    def __init__(self, max_angle: float = 22.0, alpha: float = 0.4, beta: float = 0.6, threshold: float = 0.55):
         self.detector = FaceMeshDetector()
+        self.head_pose = HeadPoseEstimator(max_angle=max_angle)
+        self.eye_scorer = EyeBehaviourScorer()
+        self.alpha = alpha
+        self.beta = beta
+        self.threshold = threshold
 
     def process_frame(self, bgr_frame: np.ndarray) -> dict:
         landmarks = self.detector.process(bgr_frame)
-        return {"landmarks": landmarks}
+        h, w = bgr_frame.shape[:2]
+
+        out = {
+            "landmarks": landmarks,
+            "s_face": 0.0,
+            "s_eye": 0.0,
+            "raw_score": 0.0,
+            "is_focused": False,
+            "yaw": None,
+            "pitch": None,
+            "roll": None,
+            "mar": None,
+            "is_yawning": False,
+        }
+
+        if landmarks is None:
+            return out
+
+        # Head pose -> S_face, yaw/pitch/roll
+        angles = self.head_pose.estimate(landmarks, w, h)
+        if angles is not None:
+            out["yaw"], out["pitch"], out["roll"] = angles
+        out["s_face"] = self.head_pose.score(landmarks, w, h)
+
+        # Eye behaviour (EAR + gaze) -> S_eye
+        out["s_eye"] = self.eye_scorer.score(landmarks)
+
+        # Mouth open (MAR) -> yawn / sleepy: force NOT FOCUSED when mouth open
+        out["mar"] = compute_mar(landmarks)
+        out["is_yawning"] = out["mar"] > MAR_YAWN_THRESHOLD
+
+        # Fusion: alpha*S_face + beta*S_eye; if yawning (mouth open) -> not focused
+        out["raw_score"] = self.alpha * out["s_face"] + self.beta * out["s_eye"]
+        out["is_focused"] = out["raw_score"] >= self.threshold and not out["is_yawning"]
+
+        return out
 
     def close(self):
         self.detector.close()