226 lines
8.8 KiB
Python
226 lines
8.8 KiB
Python
# sites/ai_mouse/ai_mouse/_collector.py
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import random
|
|
from enum import Enum, auto
|
|
from pathlib import Path
|
|
|
|
|
|
class CollectorState(Enum):
|
|
IDLE = auto()
|
|
HOVER_A = auto()
|
|
RECORDING = auto()
|
|
|
|
|
|
class Collector:
|
|
"""Manages A→B mouse movement trace collection state and persistence.
|
|
|
|
The state-machine methods (_on_mouse_motion, _on_mouseup, _on_skip)
|
|
are public so they can be called by the web API without a display.
|
|
A/B positions are exposed as .a_pos / .b_pos attributes.
|
|
"""
|
|
|
|
POINT_RADIUS = 15 # pixels — must be inside this to hover/click
|
|
DWELL_MS = 200 # milliseconds to dwell inside A before recording starts
|
|
|
|
def __init__(
|
|
self,
|
|
count: int,
|
|
dist_min: int,
|
|
dist_max: int,
|
|
output_path: Path,
|
|
screen_size: tuple[int, int] = (800, 600),
|
|
):
|
|
self.count = count
|
|
self.dist_min = dist_min
|
|
self.dist_max = dist_max
|
|
self.output_path = Path(output_path)
|
|
self.screen_w, self.screen_h = screen_size
|
|
|
|
self.collected = 0
|
|
self.state = CollectorState.IDLE
|
|
self._buffer: list[dict] = []
|
|
self._hover_enter_t: int = 0
|
|
self._record_start_t: int = 0
|
|
|
|
self.a_pos, self.b_pos = self._new_ab()
|
|
|
|
# ------------------------------------------------------------------
|
|
# State machine (called by web API)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _on_mouse_motion(self, mx: int, my: int, t: int) -> None:
|
|
"""Handle a MOUSEMOTION event at pixel (mx, my), time t ms from start."""
|
|
if self.state == CollectorState.IDLE:
|
|
if self._inside(mx, my, self.a_pos):
|
|
self.state = CollectorState.HOVER_A
|
|
self._hover_enter_t = t
|
|
|
|
elif self.state == CollectorState.HOVER_A:
|
|
if not self._inside(mx, my, self.a_pos):
|
|
self.state = CollectorState.IDLE
|
|
elif t - self._hover_enter_t >= self.DWELL_MS:
|
|
self.state = CollectorState.RECORDING
|
|
self._record_start_t = t
|
|
self._buffer = [{"type": "move", "x": mx, "y": my, "t": 0}]
|
|
|
|
elif self.state == CollectorState.RECORDING:
|
|
rel_t = t - self._record_start_t
|
|
self._buffer.append({"type": "move", "x": mx, "y": my, "t": rel_t})
|
|
|
|
def _on_mousedown(self, mx: int, my: int, t: int) -> None:
|
|
"""Handle a MOUSEBUTTONDOWN event."""
|
|
if self.state == CollectorState.RECORDING:
|
|
rel_t = t - self._record_start_t
|
|
self._buffer.append({"type": "down", "x": mx, "y": my, "t": rel_t})
|
|
|
|
def _on_mouseup(self, mx: int, my: int, t: int) -> None:
|
|
"""Handle a MOUSEBUTTONUP event."""
|
|
if self.state == CollectorState.RECORDING:
|
|
rel_t = t - self._record_start_t
|
|
self._buffer.append({"type": "up", "x": mx, "y": my, "t": rel_t})
|
|
|
|
if self._inside(mx, my, self.b_pos):
|
|
self._save_trace()
|
|
self.collected += 1
|
|
if self.collected < self.count:
|
|
self.a_pos, self.b_pos = self._new_ab()
|
|
self.state = CollectorState.IDLE
|
|
else:
|
|
# Click outside B — discard buffer and regenerate
|
|
self._on_skip()
|
|
|
|
def _on_skip(self) -> None:
|
|
"""Handle ESC/skip — discard current buffer, regenerate A/B."""
|
|
self._buffer = []
|
|
self.state = CollectorState.IDLE
|
|
self.a_pos, self.b_pos = self._new_ab()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Persistence
|
|
# ------------------------------------------------------------------
|
|
|
|
def _save_trace(self) -> None:
|
|
dist = self._dist(self.a_pos, self.b_pos)
|
|
angle = math.degrees(
|
|
math.atan2(
|
|
self.b_pos[1] - self.a_pos[1],
|
|
self.b_pos[0] - self.a_pos[0],
|
|
)
|
|
)
|
|
trace = {
|
|
"meta": {
|
|
"start": list(self.a_pos),
|
|
"end": list(self.b_pos),
|
|
"dist": round(dist),
|
|
"angle": round(angle, 1),
|
|
},
|
|
"events": list(self._buffer),
|
|
}
|
|
self.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self.output_path.open("a", encoding="utf-8") as f:
|
|
f.write(json.dumps(trace, ensure_ascii=False) + "\n")
|
|
self._buffer = []
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _inside(self, mx: int, my: int, pos: tuple[int, int]) -> bool:
|
|
return self._dist((mx, my), pos) <= self.POINT_RADIUS
|
|
|
|
@staticmethod
|
|
def _dist(a: tuple[int, int], b: tuple[int, int]) -> float:
|
|
return math.hypot(a[0] - b[0], a[1] - b[1])
|
|
|
|
def _new_ab(self) -> tuple[tuple[int, int], tuple[int, int]]:
|
|
"""Generate a new random A→B pair within distance constraints.
|
|
|
|
Strategy: clamp dist_min/dist_max to the canvas diagonal. When the
|
|
required distance is large relative to the canvas, bias A towards edges
|
|
and corners so that long-distance B positions become reachable. The
|
|
fallback randomly picks from the four corner pairs with jitter to ensure
|
|
variety even in the degenerate case.
|
|
"""
|
|
margin = self.POINT_RADIUS + 5
|
|
x_lo, x_hi = margin, self.screen_w - margin
|
|
y_lo, y_hi = margin, self.screen_h - margin
|
|
w_inner, h_inner = x_hi - x_lo, y_hi - y_lo
|
|
max_possible = int(math.hypot(w_inner, h_inner))
|
|
eff_max = min(self.dist_max, max_possible)
|
|
eff_min = min(self.dist_min, eff_max)
|
|
|
|
# Determine how "tight" the distance requirement is relative to canvas.
|
|
# When ratio > 0.7, purely random A rarely works — bias towards edges.
|
|
tightness = eff_min / max_possible if max_possible > 0 else 1.0
|
|
|
|
for _ in range(500):
|
|
if tightness > 0.7:
|
|
# Bias A towards edges/corners: pick from a ring near the border
|
|
side = random.choice(["top", "bottom", "left", "right"])
|
|
edge_band = max(int(w_inner * 0.15), 1)
|
|
if side == "top":
|
|
ax = random.randint(x_lo, x_hi)
|
|
ay = random.randint(y_lo, y_lo + edge_band)
|
|
elif side == "bottom":
|
|
ax = random.randint(x_lo, x_hi)
|
|
ay = random.randint(y_hi - edge_band, y_hi)
|
|
elif side == "left":
|
|
ax = random.randint(x_lo, x_lo + edge_band)
|
|
ay = random.randint(y_lo, y_hi)
|
|
else:
|
|
ax = random.randint(x_hi - edge_band, x_hi)
|
|
ay = random.randint(y_lo, y_hi)
|
|
else:
|
|
ax = random.randint(x_lo, x_hi)
|
|
ay = random.randint(y_lo, y_hi)
|
|
|
|
# Compute the farthest reachable distance from (ax, ay) within bounds
|
|
reach = max(
|
|
math.hypot(ax - x_lo, ay - y_lo),
|
|
math.hypot(ax - x_hi, ay - y_lo),
|
|
math.hypot(ax - x_lo, ay - y_hi),
|
|
math.hypot(ax - x_hi, ay - y_hi),
|
|
)
|
|
if reach < eff_min:
|
|
continue
|
|
|
|
local_max = min(eff_max, int(reach))
|
|
|
|
# Try several angles from this A
|
|
for _ in range(30):
|
|
angle = random.uniform(0, 2 * math.pi)
|
|
dist = random.randint(eff_min, local_max)
|
|
bx = int(ax + dist * math.cos(angle))
|
|
by = int(ay + dist * math.sin(angle))
|
|
if x_lo <= bx <= x_hi and y_lo <= by <= y_hi:
|
|
return (ax, ay), (bx, by)
|
|
|
|
# Fallback: pick a random corner pair with jitter for variety
|
|
corners = [(x_lo, y_lo), (x_hi, y_lo), (x_lo, y_hi), (x_hi, y_hi)]
|
|
pairs = [(corners[i], corners[j])
|
|
for i in range(4) for j in range(i + 1, 4)
|
|
if self._dist(corners[i], corners[j]) >= eff_min]
|
|
if not pairs:
|
|
# All pairs too short — pick the longest pair
|
|
pairs = [(corners[i], corners[j])
|
|
for i in range(4) for j in range(i + 1, 4)]
|
|
pairs.sort(key=lambda p: self._dist(p[0], p[1]), reverse=True)
|
|
pairs = pairs[:1]
|
|
|
|
ca, cb = random.choice(pairs)
|
|
# Add jitter so it's not identical each time
|
|
jitter = max(margin, int(min(w_inner, h_inner) * 0.08))
|
|
ax = ca[0] + random.randint(-jitter, jitter)
|
|
ay = ca[1] + random.randint(-jitter, jitter)
|
|
bx = cb[0] + random.randint(-jitter, jitter)
|
|
by = cb[1] + random.randint(-jitter, jitter)
|
|
# Clamp back into bounds
|
|
ax = max(x_lo, min(x_hi, ax))
|
|
ay = max(y_lo, min(y_hi, ay))
|
|
bx = max(x_lo, min(x_hi, bx))
|
|
by = max(y_lo, min(y_hi, by))
|
|
return (ax, ay), (bx, by)
|