analytics/heatmap.py - master - volkovss/eye_tracker

analytics/heatmap.py
124 строки5 KB
volkovss
Initial commit: Eye Tracker v1.0
03 июн 2026, 13:18
03 июн 2026, 13:18f9d57a2
import numpy as np
import cv2
from typing import List, Optional


def generate_heatmap(
    gaze_points:          List[dict],
    screen_w:             int,
    screen_h:             int,
    sigma:                int   = 0,
    colormap:             int   = cv2.COLORMAP_JET,
    background:           Optional[np.ndarray] = None,
    alpha:                float = 0.65,
    confidence_threshold: float = 0.25,
) -> np.ndarray:
    if not gaze_points:
        return np.zeros((screen_h, screen_w, 3), dtype=np.uint8)

    # адаптивная сигма — 4% от меньшей стороны экрана
    if sigma <= 0:
        sigma = max(30, int(min(screen_w, screen_h) * 0.04))

    heat = np.zeros((screen_h, screen_w), dtype=np.float32)

    for p in gaze_points:
        if p.get("confidence", 0) < confidence_threshold:
            continue
        x = int(np.clip(round(p["x_px"]), 0, screen_w - 1))
        y = int(np.clip(round(p["y_px"]), 0, screen_h - 1))
        heat[y, x] += max(float(p.get("confidence", 1.0)), 0.5)

    if heat.max() == 0:
        if background is not None:
            return cv2.resize(background.copy(), (screen_w, screen_h))
        return np.zeros((screen_h, screen_w, 3), dtype=np.uint8)

    ksize = sigma * 6 + 1 if (sigma * 6 + 1) % 2 == 1 else sigma * 6 + 2
    heat_big  = cv2.GaussianBlur(heat, (ksize, ksize), sigma)
    heat_med  = cv2.GaussianBlur(heat, (0, 0), sigma // 2 + 1)
    heat_combined = heat_big * 0.7 + heat_med * 0.3

    # нормализация
    heat_norm = heat_combined / heat_combined.max()

    # степенное масштабирование — подавляем слабые области,
    # усиливаем зоны концентрации
    heat_pow = np.power(heat_norm, 0.6)
    heat_u8  = (heat_pow * 255).astype(np.uint8)

    colored = cv2.applyColorMap(heat_u8, colormap)

    # маска прозрачности — показываем только где реально были точки
    alpha_mask = heat_pow[..., np.newaxis]

    if background is not None:
        bg = cv2.resize(background, (screen_w, screen_h)).astype(np.float32)
        out = (
            bg * (1.0 - alpha_mask * alpha) +
            colored.astype(np.float32) * alpha_mask * alpha
        ).astype(np.uint8)
        return out

    # без фона: тёмный фон + цветная карта
    dark = np.full((screen_h, screen_w, 3), 20, dtype=np.float32)
    out  = (dark * (1.0 - alpha_mask) + colored.astype(np.float32) * alpha_mask).astype(np.uint8)
    return out


def generate_scanpath(
    gaze_points: List[dict],
    screen_w:    int,
    screen_h:    int,
    background:  Optional[np.ndarray] = None,
    fixations=None,
    max_points:  int = 300,
) -> np.ndarray:
    if background is not None:
        canvas = cv2.resize(background.copy(), (screen_w, screen_h))
    else:
        canvas = np.full((screen_h, screen_w, 3), 20, dtype=np.uint8)

    valid = [
        (int(np.clip(p["x_px"], 0, screen_w - 1)),
         int(np.clip(p["y_px"], 0, screen_h - 1)))
        for p in gaze_points
        if p.get("confidence", 0) > 0.25
    ]

    if not valid:
        return canvas

    # прореживание: если точек много, берём равномерную выборку
    if len(valid) > max_points:
        step  = len(valid) / max_points
        valid = [valid[int(i * step)] for i in range(max_points)]

    # линии траектории с градиентом синий -> красный
    for i in range(1, len(valid)):
        t     = i / max(len(valid) - 1, 1)
        color = (int(255 * (1 - t)), 60, int(255 * t))
        cv2.line(canvas, valid[i - 1], valid[i], color, 2)

    # точки на каждой позиции
    for i, pt in enumerate(valid):
        t     = i / max(len(valid) - 1, 1)
        color = (int(255 * (1 - t)), 60, int(255 * t))
        cv2.circle(canvas, pt, 3, color, -1)

    # кружки фиксаций поверх траектории
    if fixations:
        for idx, fix in enumerate(fixations):
            # радиус пропорционален длительности, но ограничен
            r  = min(max(8, int(fix.duration_ms / 40)), 40)
            cx = int(np.clip(fix.x, 0, screen_w - 1))
            cy = int(np.clip(fix.y, 0, screen_h - 1))
            cv2.circle(canvas, (cx, cy), r, (0, 230, 230), 2)
            cv2.circle(canvas, (cx, cy), 4, (255, 230, 0), -1)
            # номер фиксации
            cv2.putText(canvas, str(idx + 1),
                        (cx + r + 2, cy + 4),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4,
                        (200, 200, 200), 1)

    return canvas