111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
# screen_capturer.py
|
|
"""
|
|
Capture d'écran pour agent_v0.
|
|
|
|
v0+ :
|
|
- capture de l'écran principal en PNG
|
|
- mode "full" : écran complet
|
|
- mode "crop" : zone centrée autour d'une position (curseur)
|
|
- stockage dans sessions/<session_id>/shots/shot_XXXX.png
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from typing import Tuple, Optional
|
|
|
|
from mss import mss
|
|
from mss.tools import to_png
|
|
|
|
from config import SESSIONS_ROOT
|
|
from raw_session import RawSession
|
|
|
|
|
|
class ScreenCapturer:
|
|
"""
|
|
Gère la capture d'écran pour une session donnée.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
session: RawSession,
|
|
base_dir: str = SESSIONS_ROOT,
|
|
screenshot_mode: str = "full", # "full" | "crop"
|
|
crop_width: int = 800,
|
|
crop_height: int = 600,
|
|
) -> None:
|
|
self.session = session
|
|
self.base_dir = base_dir
|
|
self._counter = 0
|
|
|
|
self.screenshot_mode = screenshot_mode
|
|
self.crop_width = crop_width
|
|
self.crop_height = crop_height
|
|
|
|
# Résolution écran pour clamp la zone de crop
|
|
self.screen_width, self.screen_height = self.detect_primary_resolution()
|
|
|
|
def _get_session_shots_dir(self) -> str:
|
|
session_dir = os.path.join(self.base_dir, self.session.session_id)
|
|
shots_dir = os.path.join(session_dir, "shots")
|
|
os.makedirs(shots_dir, exist_ok=True)
|
|
return shots_dir
|
|
|
|
def capture(self, focus_pos: Optional[Tuple[int, int]] = None) -> Tuple[str, str]:
|
|
"""
|
|
Capture l'écran :
|
|
- si screenshot_mode == "crop" et focus_pos fourni -> zone centrée autour de focus_pos
|
|
- sinon -> plein écran
|
|
|
|
Retourne (screenshot_id, relative_path)
|
|
"""
|
|
self._counter += 1
|
|
screenshot_id = f"shot_{self._counter:04d}"
|
|
shots_dir = self._get_session_shots_dir()
|
|
filename = f"{screenshot_id}.png"
|
|
file_path = os.path.join(shots_dir, filename)
|
|
relative_path = os.path.join("shots", filename)
|
|
|
|
with mss() as sct:
|
|
if self.screenshot_mode == "crop" and focus_pos is not None:
|
|
cx, cy = focus_pos
|
|
w = min(self.crop_width, self.screen_width)
|
|
h = min(self.crop_height, self.screen_height)
|
|
|
|
left = max(0, cx - w // 2)
|
|
top = max(0, cy - h // 2)
|
|
|
|
if left + w > self.screen_width:
|
|
left = self.screen_width - w
|
|
if top + h > self.screen_height:
|
|
top = self.screen_height - h
|
|
|
|
region = {"left": int(left), "top": int(top), "width": int(w), "height": int(h)}
|
|
img = sct.grab(region)
|
|
else:
|
|
monitor = sct.monitors[0] # écran principal complet
|
|
img = sct.grab(monitor)
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(to_png(img.rgb, img.size))
|
|
|
|
captured_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
self.session.add_screenshot(
|
|
screenshot_id=screenshot_id,
|
|
relative_path=relative_path,
|
|
captured_at=captured_at,
|
|
)
|
|
|
|
return screenshot_id, relative_path
|
|
|
|
@staticmethod
|
|
def detect_primary_resolution() -> Tuple[int, int]:
|
|
"""
|
|
Renvoie (width, height) de l'écran principal.
|
|
Utile pour remplir RawSession.environment.screen.primary_resolution.
|
|
"""
|
|
with mss() as sct:
|
|
mon = sct.monitors[0]
|
|
return mon["width"], mon["height"]
|