main.py

"""
AttenSense - classroom engagement monitor
Local CV pipeline (YOLO26 + MediaPipe) + FastAPI backend + SQLite storage.

Setup:
    pip install -r requirements.txt

Run:
    uvicorn main:app --reload

Then open:
    http://localhost:8000
"""
import os
import sqlite3
import threading
import time
import urllib.request
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np
from fastapi import FastAPI
from fastapi.responses import StreamingResponse, FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from ultralytics import YOLO
import mediapipe as mp
from mediapipe.tasks import python as mp_tasks
from mediapipe.tasks.python import vision as mp_vision

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).parent
SNAPSHOT_DIR = BASE_DIR / "snapshots"
SNAPSHOT_DIR.mkdir(exist_ok=True)
DB_PATH = BASE_DIR / "attensense.db"

# ---------------------------------------------------------------------------
# Database (no names, no student accounts - just timestamped events)
# ---------------------------------------------------------------------------
def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT NOT NULL,
            reason TEXT NOT NULL,
            image_path TEXT NOT NULL
        )
    """)
    conn.commit()
    conn.close()


def save_event(reason: str, frame) -> None:
    ts = datetime.now()
    filename = f"{ts.strftime('%Y%m%d_%H%M%S_%f')}.jpg"
    cv2.imwrite(str(SNAPSHOT_DIR / filename), frame)
    conn = sqlite3.connect(DB_PATH)
    conn.execute(
        "INSERT INTO events (timestamp, reason, image_path) VALUES (?, ?, ?)",
        (ts.strftime("%I:%M %p"), reason, filename),
    )
    conn.commit()
    conn.close()


def get_events(limit: int = 100):
    conn = sqlite3.connect(DB_PATH)
    rows = conn.execute(
        "SELECT timestamp, reason, image_path FROM events ORDER BY id DESC LIMIT ?",
        (limit,),
    ).fetchall()
    conn.close()
    return [{"timestamp": r[0], "reason": r[1], "image_path": r[2]} for r in rows]


init_db()

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------
# YOLO26n: nano variant, fastest CPU inference in the YOLO26 family - the
# right pick for a laptop webcam demo. Swap to yolo26s.pt if you have a GPU
# and want more accuracy. Weights auto-download on first run.
yolo = YOLO("yolo26n.pt")
PERSON_CLASS = 0
PHONE_CLASS = 67  # "cell phone" in the COCO class list

# MediaPipe dropped the old mp.solutions.face_mesh API in favor of the
# MediaPipe Tasks API. FaceLandmarker is the current replacement; unlike
# YOLO it doesn't auto-download its model, so we fetch the .task file
# ourselves on first run (a few MB, cached in BASE_DIR after that).
FACE_LANDMARKER_MODEL_PATH = BASE_DIR / "face_landmarker.task"
FACE_LANDMARKER_MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/face_landmarker/"
    "face_landmarker/float16/latest/face_landmarker.task"
)

if not FACE_LANDMARKER_MODEL_PATH.exists():
    print("Downloading MediaPipe face landmarker model (first run only)...")
    urllib.request.urlretrieve(FACE_LANDMARKER_MODEL_URL, FACE_LANDMARKER_MODEL_PATH)

face_landmarker = mp_vision.FaceLandmarker.create_from_options(
    mp_vision.FaceLandmarkerOptions(
        base_options=mp_tasks.BaseOptions(model_asset_path=str(FACE_LANDMARKER_MODEL_PATH)),
        running_mode=mp_vision.RunningMode.IMAGE,
        num_faces=10,
        min_face_detection_confidence=0.5,
        min_face_presence_confidence=0.5,
        min_tracking_confidence=0.5,
    )
)

# Eye landmark indices (MediaPipe's 468-point face mesh)
LEFT_EYE = [362, 385, 387, 263, 373, 380]
RIGHT_EYE = [33, 160, 158, 133, 153, 144]
EAR_THRESHOLD = 0.21     # below this = eyes closing
DROWSY_FRAMES = 15       # consecutive low-EAR frames (~0.5s at 30fps) to confirm

# ---------------------------------------------------------------------------
# Shared state between the capture thread and the API
# ---------------------------------------------------------------------------
lock = threading.Lock()
latest_jpeg = None
stats = {"engaged": 0, "distracted": 0, "drowsy": 0, "running": False}
drowsy_counters = {}   # face index -> consecutive low-EAR frame count
cooldowns = {}          # reason -> last-saved timestamp (avoid snapshot spam)


def eye_aspect_ratio(landmarks, eye_idx, w, h):
    p1, p2, p3, p4, p5, p6 = [(landmarks[i].x * w, landmarks[i].y * h) for i in eye_idx]
    vertical = np.linalg.norm(np.array(p2) - np.array(p6)) + np.linalg.norm(np.array(p3) - np.array(p5))
    horizontal = np.linalg.norm(np.array(p1) - np.array(p4))
    return vertical / (2.0 * horizontal + 1e-6)


def maybe_save(reason: str, frame, cooldown_seconds: int = 10):
    now = time.time()
    if now - cooldowns.get(reason, 0) > cooldown_seconds:
        cooldowns[reason] = now
        save_event(reason, frame)


def open_camera():
    """Open the webcam.

    On Linux, cv2.VideoCapture(index) can be slow or pick the wrong backend
    if left to auto-detect (especially with virtual devices from apps like
    OBS, or a laptop having more than one /dev/video* node). We try the
    V4L2 backend explicitly first, since it's the one that actually talks
    to /dev/video*, then fall back to the default backend if that fails.

    The device index defaults to 0 but can be overridden with the
    CAMERA_INDEX env var, e.g.:  CAMERA_INDEX=1 uvicorn main:app --reload
    Run `v4l2-ctl --list-devices` (from the v4l-utils package) to find the
    right index if 0 isn't your webcam.
    """
    index = int(os.environ.get("CAMERA_INDEX", "0"))

    cap = cv2.VideoCapture(index, cv2.CAP_V4L2)
    if not cap.isOpened():
        cap.release()
        cap = cv2.VideoCapture(index)  # let OpenCV auto-pick a backend
    return cap


def capture_loop():
    global latest_jpeg
    cap = open_camera()
    if not cap.isOpened():
        print(
            "Could not open webcam at CAMERA_INDEX="
            f"{os.environ.get('CAMERA_INDEX', '0')}. On Linux, check:\n"
            "  - the device exists: ls /dev/video*\n"
            "  - your user has access: groups $USER (should include 'video')\n"
            "  - the right index: v4l2-ctl --list-devices"
        )
        stats["running"] = False
        return

    while stats["running"]:
        ok, frame = cap.read()
        if not ok:
            continue

        h, w = frame.shape[:2]
        engaged = distracted = drowsy = 0

        # --- YOLO26: people + phones in the frame ---
        results = yolo.predict(frame, verbose=False, classes=[PERSON_CLASS, PHONE_CLASS])[0]
        phone_boxes = []
        for box in results.boxes:
            cls = int(box.cls[0])
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            if cls == PHONE_CLASS:
                phone_boxes.append((x1, y1, x2, y2))
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 165, 255), 2)

        # --- MediaPipe: per-face drowsiness + phone-proximity distraction ---
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
        mesh_results = face_landmarker.detect(mp_image)

        seen_faces = set()
        if mesh_results.face_landmarks:
            for idx, lm in enumerate(mesh_results.face_landmarks):
                seen_faces.add(idx)
                ear = (eye_aspect_ratio(lm, LEFT_EYE, w, h) + eye_aspect_ratio(lm, RIGHT_EYE, w, h)) / 2

                xs = [p.x * w for p in lm]
                ys = [p.y * h for p in lm]
                x1, y1, x2, y2 = int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))

                is_drowsy = False
                if ear < EAR_THRESHOLD:
                    drowsy_counters[idx] = drowsy_counters.get(idx, 0) + 1
                    is_drowsy = drowsy_counters[idx] >= DROWSY_FRAMES
                else:
                    drowsy_counters[idx] = 0

                is_distracted = any(
                    px1 < x2 and px2 > x1 and py1 < y2 + 100 and py2 > y1
                    for px1, py1, px2, py2 in phone_boxes
                )

                if is_drowsy:
                    drowsy += 1
                    color = (0, 0, 255)
                    maybe_save("Drowsiness", frame)
                elif is_distracted:
                    distracted += 1
                    color = (0, 165, 255)
                    maybe_save("Distraction", frame)
                else:
                    engaged += 1
                    color = (0, 200, 0)

                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

        for idx in list(drowsy_counters):
            if idx not in seen_faces:
                drowsy_counters.pop(idx, None)

        with lock:
            stats["engaged"], stats["distracted"], stats["drowsy"] = engaged, distracted, drowsy
            ok, buf = cv2.imencode(".jpg", frame)
            if ok:
                latest_jpeg = buf.tobytes()

    cap.release()


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------
app = FastAPI(title="AttenSense")
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
app.mount("/snapshots", StaticFiles(directory=str(SNAPSHOT_DIR)), name="snapshots")


@app.get("/")
def index():
    return FileResponse(str(BASE_DIR / "static" / "index.html"))


@app.post("/api/start")
def start_class():
    if not stats["running"]:
        stats["running"] = True
        threading.Thread(target=capture_loop, daemon=True).start()
    return {"running": True}


@app.post("/api/stop")
def stop_class():
    stats["running"] = False
    return {"running": False}


@app.get("/api/stats")
def get_stats():
    with lock:
        return JSONResponse({**stats})


@app.get("/api/disengaged")
def disengaged():
    return JSONResponse(get_events())


def mjpeg_generator():
    while True:
        with lock:
            frame = latest_jpeg
        if frame is not None:
            yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + frame + b"\r\n")
        time.sleep(0.03)


@app.get("/video_feed")
def video_feed():
    return StreamingResponse(mjpeg_generator(), media_type="multipart/x-mixed-replace; boundary=frame")