Initial commit

screen_capturer.py

@@ -0,0 +1,110 @@
+# 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"]