├── controls ├── keyboard │ └── __init__.py ├── __init__.py └── mouse │ ├── base.py │ ├── pyautogui_mouse.py │ ├── pydirectinput_mouse.py │ ├── pynput_mouse.py │ ├── __init__.py │ └── win32_mouse.py ├── .idea ├── .gitignore ├── vcs.xml ├── inspectionProfiles │ ├── profiles_settings.xml │ └── Project_Default.xml ├── modules.xml ├── misc.xml └── aicapturebase.iml ├── streaming ├── __init__.py └── client.py ├── requirements.txt ├── exceptions.py ├── utils ├── timing.py ├── __init__.py ├── fps.py ├── win32.py ├── nms.py ├── benchmark.py ├── windmouse.py └── cv.py ├── grabbers ├── base.py ├── mss_grabber.py ├── d3dshot_grabber.py ├── dxcam_grabber.py ├── screengear_grabber.py ├── dxcam_capture_grabber.py ├── win32_grabber.py ├── __init__.py └── obs_vc_grabber.py ├── README.md ├── LICENSE ├── pyproject.toml ├── config.py ├── .gitignore └── example.py /controls/keyboard/__init__.py: -------------------------------------------------------------------------------- 1 | # Keyboard controls placeholder 2 | -------------------------------------------------------------------------------- /.idea/.gitignore: -------------------------------------------------------------------------------- 1 | # Default ignored files 2 | /shelf/ 3 | /workspace.xml 4 | -------------------------------------------------------------------------------- /streaming/__init__.py: -------------------------------------------------------------------------------- 1 | from .client import StreamClient, run_viewer 2 | 3 | __all__ = ["StreamClient", "run_viewer"] 4 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | dxcam 2 | keyboard 3 | mss 4 | numpy 5 | opencv-python 6 | PyAutoGUI 7 | PyDirectInput 8 | pynput 9 | pywin32 10 | pygrabber 11 | lz4 12 | -------------------------------------------------------------------------------- /.idea/vcs.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /controls/__init__.py: -------------------------------------------------------------------------------- 1 | from .mouse import BaseMouseControls, get_mouse_controls, list_mouse_controls 2 | 3 | __all__ = [ 4 | "BaseMouseControls", 5 | "get_mouse_controls", 6 | "list_mouse_controls", 7 | ] 8 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/profiles_settings.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 6 | -------------------------------------------------------------------------------- /.idea/modules.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.idea/misc.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 6 | 7 | -------------------------------------------------------------------------------- /exceptions.py: -------------------------------------------------------------------------------- 1 | class AICaptureError(Exception): 2 | pass 3 | 4 | 5 | class GrabberError(AICaptureError): 6 | pass 7 | 8 | 9 | class GrabberInitError(GrabberError): 10 | pass 11 | 12 | 13 | class DeviceNotFoundError(GrabberError): 14 | pass 15 | 16 | 17 | class CaptureError(GrabberError): 18 | pass 19 | 20 | 21 | class ControlsError(AICaptureError): 22 | pass 23 | -------------------------------------------------------------------------------- /.idea/aicapturebase.iml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | -------------------------------------------------------------------------------- /utils/timing.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | def precise_sleep(duration: float, precision: float = 0.0005) -> None: 5 | if duration <= 0: 6 | return 7 | 8 | end = time.perf_counter() + duration 9 | 10 | if duration > precision * 2: 11 | time.sleep(duration - precision) 12 | 13 | while time.perf_counter() < end: 14 | pass 15 | 16 | 17 | def busy_sleep(duration: float) -> None: 18 | if duration <= 0: 19 | return 20 | 21 | end = time.perf_counter() + duration 22 | while time.perf_counter() < end: 23 | pass 24 | -------------------------------------------------------------------------------- /grabbers/base.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import Dict, Optional, Any 3 | 4 | import numpy as np 5 | 6 | 7 | class BaseGrabber(ABC): 8 | 9 | _type: str = "base" 10 | 11 | @property 12 | def type(self) -> str: 13 | return self._type 14 | 15 | def initialize(self, **kwargs: Any) -> None: 16 | pass 17 | 18 | @abstractmethod 19 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 20 | pass 21 | 22 | def cleanup(self) -> None: 23 | pass 24 | 25 | def __enter__(self): 26 | return self 27 | 28 | def __exit__(self, exc_type, exc_val, exc_tb): 29 | self.cleanup() 30 | return False 31 | -------------------------------------------------------------------------------- /grabbers/mss_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional 2 | 3 | import mss 4 | import numpy as np 5 | 6 | from .base import BaseGrabber 7 | 8 | 9 | class MSSGrabber(BaseGrabber): 10 | 11 | _type = "mss" 12 | 13 | def __init__(self): 14 | self._sct: Optional[mss.mss] = None 15 | 16 | def initialize(self, **kwargs) -> None: 17 | self._sct = mss.mss() 18 | 19 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 20 | if self._sct is None: 21 | self.initialize() 22 | return np.array(self._sct.grab(grab_area)) 23 | 24 | def cleanup(self) -> None: 25 | if self._sct is not None: 26 | self._sct.close() 27 | self._sct = None 28 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AI Capture Basic Template 2 | Basic code template to start your new AI detection/processing project with real-time image capturing (i.e. of a game window). 3 | 4 | ## Supports different capture backends: 5 | - d3dshot 6 | - dxcam 7 | - dxcam (capture mode) 8 | - mss 9 | - screengear 10 | - win32 11 | - obs _(via virtual camera)_ 12 | 13 | ## Mouse control libraries included 14 | - pyAutoGui 15 | - pyDirectInput 16 | - pynput 17 | - win32 18 | 19 | ## Other 20 | - Windmouse algorithm included (realistic mouse movement simulation) 21 | - Benchmarks class 22 | - NMS stuff 23 | - CV2 utils 24 | - etc utils 25 | 26 | ## Not done yet 27 | There's a 2 and more PCs streaming code included (client only). 28 | That lets you run application/game on first PC, and do inference etc on another. 29 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .fps import FPSCounter, FrameTimer 2 | from .timing import precise_sleep, busy_sleep 3 | from .benchmark import Benchmark 4 | from .cv import ( 5 | combine_bboxes, 6 | xywh_to_xyxy, 7 | xyxy_to_xywh, 8 | calc_iou, 9 | boxes_intersect, 10 | merge_overlapping_boxes, 11 | point_offset, 12 | round_to_multiple, 13 | ) 14 | from .nms import non_max_suppression 15 | from .windmouse import wind_mouse 16 | 17 | __all__ = [ 18 | "FPSCounter", 19 | "FrameTimer", 20 | "precise_sleep", 21 | "busy_sleep", 22 | "Benchmark", 23 | "combine_bboxes", 24 | "xywh_to_xyxy", 25 | "xyxy_to_xywh", 26 | "calc_iou", 27 | "boxes_intersect", 28 | "merge_overlapping_boxes", 29 | "point_offset", 30 | "round_to_multiple", 31 | "non_max_suppression", 32 | "wind_mouse", 33 | ] 34 | -------------------------------------------------------------------------------- /grabbers/d3dshot_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Any 2 | 3 | import numpy as np 4 | 5 | from .base import BaseGrabber 6 | 7 | 8 | class D3DShotGrabber(BaseGrabber): 9 | 10 | _type = "d3dshot" 11 | 12 | def __init__(self): 13 | self._d3d = None 14 | 15 | def initialize(self, **kwargs: Any) -> None: 16 | import d3dshot 17 | self._d3d = d3dshot.create(capture_output="numpy") 18 | 19 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 20 | if self._d3d is None: 21 | self.initialize() 22 | 23 | region = ( 24 | grab_area["left"], 25 | grab_area["top"], 26 | grab_area["left"] + grab_area["width"], 27 | grab_area["top"] + grab_area["height"], 28 | ) 29 | return self._d3d.screenshot(region=region) 30 | 31 | def cleanup(self) -> None: 32 | self._d3d = None 33 | -------------------------------------------------------------------------------- /controls/mouse/base.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import Tuple 3 | 4 | 5 | class BaseMouseControls(ABC): 6 | 7 | _type: str = "base" 8 | 9 | @property 10 | def type(self) -> str: 11 | return self._type 12 | 13 | @abstractmethod 14 | def move(self, x: int, y: int) -> None: 15 | pass 16 | 17 | @abstractmethod 18 | def move_relative(self, dx: int, dy: int) -> None: 19 | pass 20 | 21 | @abstractmethod 22 | def click(self, button: str = "left") -> None: 23 | pass 24 | 25 | @abstractmethod 26 | def get_position(self) -> Tuple[int, int]: 27 | pass 28 | 29 | def press(self, button: str = "left") -> None: 30 | pass 31 | 32 | def release(self, button: str = "left") -> None: 33 | pass 34 | 35 | def double_click(self, button: str = "left") -> None: 36 | self.click(button) 37 | self.click(button) 38 | -------------------------------------------------------------------------------- /grabbers/dxcam_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Any 2 | 3 | import numpy as np 4 | 5 | from .base import BaseGrabber 6 | 7 | 8 | class DXCamGrabber(BaseGrabber): 9 | 10 | _type = "dxcam" 11 | 12 | def __init__(self): 13 | self._camera = None 14 | 15 | def initialize(self, **kwargs: Any) -> None: 16 | import dxcam 17 | self._camera = dxcam.create() 18 | 19 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 20 | if self._camera is None: 21 | self.initialize() 22 | 23 | region = ( 24 | grab_area["left"], 25 | grab_area["top"], 26 | grab_area["left"] + grab_area["width"], 27 | grab_area["top"] + grab_area["height"], 28 | ) 29 | return self._camera.grab(region=region) 30 | 31 | def cleanup(self) -> None: 32 | if self._camera is not None: 33 | del self._camera 34 | self._camera = None 35 | -------------------------------------------------------------------------------- /controls/mouse/pyautogui_mouse.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import pyautogui 4 | 5 | from .base import BaseMouseControls 6 | 7 | pyautogui.MINIMUM_DURATION = 0 8 | pyautogui.MINIMUM_SLEEP = 0 9 | pyautogui.PAUSE = 0 10 | pyautogui.FAILSAFE = False 11 | 12 | 13 | class PyAutoGUIMouseControls(BaseMouseControls): 14 | 15 | _type = "pyautogui" 16 | 17 | def move(self, x: int, y: int) -> None: 18 | pyautogui.moveTo(x, y) 19 | 20 | def move_relative(self, dx: int, dy: int) -> None: 21 | pyautogui.moveRel(dx, dy) 22 | 23 | def click(self, button: str = "left") -> None: 24 | pyautogui.click(button=button) 25 | 26 | def get_position(self) -> Tuple[int, int]: 27 | pos = pyautogui.position() 28 | return (pos.x, pos.y) 29 | 30 | def press(self, button: str = "left") -> None: 31 | pyautogui.mouseDown(button=button) 32 | 33 | def release(self, button: str = "left") -> None: 34 | pyautogui.mouseUp(button=button) 35 | -------------------------------------------------------------------------------- /controls/mouse/pydirectinput_mouse.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import pydirectinput 4 | 5 | from .base import BaseMouseControls 6 | 7 | pydirectinput.MINIMUM_DURATION = 0 8 | pydirectinput.MINIMUM_SLEEP = 0 9 | pydirectinput.PAUSE = 0 10 | pydirectinput.FAILSAFE = False 11 | 12 | 13 | class PyDirectInputMouseControls(BaseMouseControls): 14 | 15 | _type = "pydirectinput" 16 | 17 | def move(self, x: int, y: int) -> None: 18 | pydirectinput.moveTo(x, y) 19 | 20 | def move_relative(self, dx: int, dy: int) -> None: 21 | pydirectinput.moveRel(dx, dy) 22 | 23 | def click(self, button: str = "left") -> None: 24 | pydirectinput.click(button=button) 25 | 26 | def get_position(self) -> Tuple[int, int]: 27 | pos = pydirectinput.position() 28 | return (pos[0], pos[1]) 29 | 30 | def press(self, button: str = "left") -> None: 31 | pydirectinput.mouseDown(button=button) 32 | 33 | def release(self, button: str = "left") -> None: 34 | pydirectinput.mouseUp(button=button) 35 | -------------------------------------------------------------------------------- /utils/fps.py: -------------------------------------------------------------------------------- 1 | import time 2 | from collections import deque 3 | from typing import Optional 4 | 5 | 6 | class FPSCounter: 7 | 8 | def __init__(self, sample_size: int = 50): 9 | self._timestamps: deque = deque(maxlen=sample_size) 10 | 11 | def __call__(self) -> float: 12 | self._timestamps.append(time.perf_counter()) 13 | 14 | if len(self._timestamps) < 2: 15 | return 0.0 16 | 17 | delta = self._timestamps[-1] - self._timestamps[0] 18 | if delta <= 0: 19 | return 0.0 20 | 21 | return (len(self._timestamps) - 1) / delta 22 | 23 | def reset(self) -> None: 24 | self._timestamps.clear() 25 | 26 | 27 | class FrameTimer: 28 | 29 | def __init__(self): 30 | self._start: Optional[float] = None 31 | 32 | def start(self) -> None: 33 | self._start = time.perf_counter() 34 | 35 | def elapsed(self) -> float: 36 | if self._start is None: 37 | return 0.0 38 | return time.perf_counter() - self._start 39 | 40 | def elapsed_ms(self) -> float: 41 | return self.elapsed() * 1000 42 | -------------------------------------------------------------------------------- /grabbers/screengear_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Any 2 | 3 | import numpy as np 4 | 5 | from .base import BaseGrabber 6 | 7 | 8 | class ScreenGearGrabber(BaseGrabber): 9 | 10 | _type = "screengear" 11 | 12 | def __init__(self): 13 | self._stream = None 14 | self._initialized = False 15 | self._logging = False 16 | 17 | def initialize(self, logging: bool = False, **kwargs: Any) -> None: 18 | self._logging = logging 19 | 20 | def _start_stream(self, grab_area: Dict[str, int]) -> None: 21 | from vidgear.gears import ScreenGear 22 | self._stream = ScreenGear(logging=self._logging, **grab_area).start() 23 | self._initialized = True 24 | 25 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 26 | if not self._initialized: 27 | self._start_stream(grab_area) 28 | return self._stream.read() 29 | 30 | def cleanup(self) -> None: 31 | if self._stream is not None: 32 | self._stream.stop() 33 | self._stream = None 34 | self._initialized = False 35 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Priler 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /utils/win32.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Optional 2 | 3 | 4 | def get_window_rect( 5 | window_title: str, 6 | border_offsets: Tuple[int, int, int, int] = (0, 0, 0, 0), 7 | ) -> Tuple[int, int, int, int]: 8 | import win32gui 9 | 10 | if not window_title: 11 | raise ValueError("Window title cannot be empty") 12 | 13 | hwnd = win32gui.FindWindow(None, window_title) 14 | if not hwnd: 15 | raise ValueError(f"Window not found: {window_title}") 16 | 17 | rect = list(win32gui.GetWindowRect(hwnd)) 18 | 19 | width = rect[2] - rect[0] 20 | height = rect[3] - rect[1] 21 | 22 | left = rect[0] + border_offsets[0] 23 | top = rect[1] + border_offsets[1] 24 | width -= border_offsets[0] + border_offsets[2] 25 | height -= border_offsets[1] + border_offsets[3] 26 | 27 | return (left, top, width, height) 28 | 29 | 30 | def find_window(title: str) -> Optional[int]: 31 | import win32gui 32 | return win32gui.FindWindow(None, title) or None 33 | 34 | 35 | def get_foreground_window_title() -> str: 36 | import win32gui 37 | hwnd = win32gui.GetForegroundWindow() 38 | return win32gui.GetWindowText(hwnd) 39 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "aicapturebase" 7 | version = "0.1.0" 8 | description = "High-performance screen capture template for AI/ML applications" 9 | readme = "README.md" 10 | requires-python = ">=3.9" 11 | license = "MIT" 12 | keywords = ["screen-capture", "ai", "computer-vision", "automation"] 13 | 14 | dependencies = [ 15 | "numpy>=1.20.0", 16 | "opencv-python>=4.5.0", 17 | "mss>=6.0.0", 18 | ] 19 | 20 | [project.optional-dependencies] 21 | win32 = [ 22 | "pywin32>=300", 23 | "dxcam>=0.0.5", 24 | "pygrabber>=0.1", 25 | ] 26 | controls = [ 27 | "keyboard>=0.13.5", 28 | "pynput>=1.7.0", 29 | "PyAutoGUI>=0.9.50", 30 | "PyDirectInput>=1.0.4", 31 | ] 32 | streaming = [ 33 | "lz4>=4.0.0", 34 | ] 35 | all = ["aicapturebase[win32,controls,streaming]"] 36 | 37 | [tool.ruff] 38 | line-length = 100 39 | target-version = "py39" 40 | 41 | [tool.ruff.lint] 42 | select = ["E", "F", "W", "I", "N", "UP", "B", "C4"] 43 | 44 | [tool.autopep8] 45 | max_line_length = 120 46 | ignore = "E501,W6" 47 | in-place = true 48 | recursive = true 49 | aggressive = 3 -------------------------------------------------------------------------------- /utils/nms.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def non_max_suppression(boxes: np.ndarray, overlap_thresh: float = 0.3) -> np.ndarray: 5 | if len(boxes) == 0: 6 | return np.array([], dtype=np.int32).reshape(0, 4) 7 | 8 | if boxes.dtype.kind == "i": 9 | boxes = boxes.astype(np.float64) 10 | 11 | pick = [] 12 | 13 | x1 = boxes[:, 0] 14 | y1 = boxes[:, 1] 15 | x2 = boxes[:, 2] 16 | y2 = boxes[:, 3] 17 | 18 | area = (x2 - x1 + 1) * (y2 - y1 + 1) 19 | idxs = np.argsort(y2) 20 | 21 | while len(idxs) > 0: 22 | last = len(idxs) - 1 23 | i = idxs[last] 24 | pick.append(i) 25 | 26 | xx1 = np.maximum(x1[i], x1[idxs[:last]]) 27 | yy1 = np.maximum(y1[i], y1[idxs[:last]]) 28 | xx2 = np.minimum(x2[i], x2[idxs[:last]]) 29 | yy2 = np.minimum(y2[i], y2[idxs[:last]]) 30 | 31 | w = np.maximum(0, xx2 - xx1 + 1) 32 | h = np.maximum(0, yy2 - yy1 + 1) 33 | 34 | overlap = (w * h) / area[idxs[:last]] 35 | 36 | idxs = np.delete( 37 | idxs, 38 | np.concatenate(([last], np.where(overlap > overlap_thresh)[0])) 39 | ) 40 | 41 | return boxes[pick].astype(np.int32) 42 | -------------------------------------------------------------------------------- /controls/mouse/pynput_mouse.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | from pynput.mouse import Button, Controller 4 | 5 | from .base import BaseMouseControls 6 | 7 | 8 | class PynputMouseControls(BaseMouseControls): 9 | 10 | _type = "pynput" 11 | 12 | def __init__(self): 13 | self._controller = Controller() 14 | 15 | def _get_button(self, button: str) -> Button: 16 | buttons = { 17 | "left": Button.left, 18 | "right": Button.right, 19 | "middle": Button.middle, 20 | } 21 | return buttons.get(button, Button.left) 22 | 23 | def move(self, x: int, y: int) -> None: 24 | self._controller.position = (x, y) 25 | 26 | def move_relative(self, dx: int, dy: int) -> None: 27 | self._controller.move(dx, dy) 28 | 29 | def click(self, button: str = "left") -> None: 30 | self._controller.click(self._get_button(button), 1) 31 | 32 | def get_position(self) -> Tuple[int, int]: 33 | return self._controller.position 34 | 35 | def press(self, button: str = "left") -> None: 36 | self._controller.press(self._get_button(button)) 37 | 38 | def release(self, button: str = "left") -> None: 39 | self._controller.release(self._get_button(button)) 40 | -------------------------------------------------------------------------------- /utils/benchmark.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import Dict, Optional, Tuple 3 | from contextlib import contextmanager 4 | 5 | 6 | class Benchmark: 7 | 8 | def __init__(self): 9 | self._marks: Dict[str, list] = {} 10 | 11 | def start(self, name: str = "default") -> float: 12 | now = time.perf_counter() 13 | self._marks[name] = [0.0, now, 0.0] 14 | return now 15 | 16 | def end(self, name: str = "default") -> Tuple[float, str]: 17 | now = time.perf_counter() 18 | 19 | if name not in self._marks: 20 | return (0.0, "0ms") 21 | 22 | self._marks[name][2] = now 23 | elapsed = now - self._marks[name][1] 24 | self._marks[name][0] = elapsed 25 | 26 | return (elapsed, f"{int(elapsed * 1000)}ms") 27 | 28 | def clear(self, name: Optional[str] = None) -> None: 29 | if name is None: 30 | self._marks.clear() 31 | elif name in self._marks: 32 | del self._marks[name] 33 | 34 | @contextmanager 35 | def measure(self, name: str = "default"): 36 | self.start(name) 37 | try: 38 | yield 39 | finally: 40 | self.end(name) 41 | 42 | def get_result(self, name: str = "default") -> Optional[float]: 43 | if name in self._marks: 44 | return self._marks[name][0] 45 | return None 46 | -------------------------------------------------------------------------------- /grabbers/dxcam_capture_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Any 2 | 3 | import numpy as np 4 | 5 | from .base import BaseGrabber 6 | from exceptions import GrabberInitError 7 | 8 | 9 | class DXCamCaptureGrabber(BaseGrabber): 10 | 11 | _type = "dxcam_capture" 12 | 13 | def __init__(self): 14 | self._camera = None 15 | self._initialized = False 16 | self._region = None 17 | 18 | def initialize(self, **kwargs: Any) -> None: 19 | import dxcam 20 | self._camera = dxcam.create() 21 | 22 | def _start_capture(self, grab_area: Dict[str, int]) -> None: 23 | if self._camera is None: 24 | self.initialize() 25 | 26 | self._region = ( 27 | grab_area["left"], 28 | grab_area["top"], 29 | grab_area["left"] + grab_area["width"], 30 | grab_area["top"] + grab_area["height"], 31 | ) 32 | 33 | self._camera.start(region=self._region) 34 | 35 | if not self._camera.is_capturing: 36 | raise GrabberInitError("DXCam failed to start capture") 37 | 38 | self._initialized = True 39 | 40 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 41 | if not self._initialized: 42 | self._start_capture(grab_area) 43 | return self._camera.get_latest_frame() 44 | 45 | def cleanup(self) -> None: 46 | if self._camera is not None: 47 | if self._camera.is_capturing: 48 | self._camera.stop() 49 | del self._camera 50 | self._camera = None 51 | self._initialized = False 52 | -------------------------------------------------------------------------------- /controls/mouse/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Type 2 | 3 | from .base import BaseMouseControls 4 | 5 | _CONTROLS: Dict[str, Type[BaseMouseControls]] = {} 6 | _LOAD_ERRORS: Dict[str, str] = {} 7 | 8 | 9 | def get_mouse_controls(name: str) -> BaseMouseControls: 10 | _ensure_controls_loaded() 11 | 12 | if name not in _CONTROLS: 13 | available = list(_CONTROLS.keys()) 14 | error_msg = f"Unknown mouse control: {name}. Available: {available}" 15 | if name in _LOAD_ERRORS: 16 | error_msg += f"\nLoad error for '{name}': {_LOAD_ERRORS[name]}" 17 | raise ValueError(error_msg) 18 | 19 | return _CONTROLS[name]() 20 | 21 | 22 | def list_mouse_controls() -> list: 23 | _ensure_controls_loaded() 24 | return list(_CONTROLS.keys()) 25 | 26 | 27 | def _try_load(name: str, module: str, class_name: str) -> None: 28 | try: 29 | mod = __import__(module, fromlist=[class_name]) 30 | _CONTROLS[name] = getattr(mod, class_name) 31 | except Exception as e: 32 | _LOAD_ERRORS[name] = f"{type(e).__name__}: {e}" 33 | 34 | 35 | def _ensure_controls_loaded(): 36 | if _CONTROLS: 37 | return 38 | 39 | _try_load("win32", "controls.mouse.win32_mouse", "Win32MouseControls") 40 | _try_load("pyautogui", "controls.mouse.pyautogui_mouse", "PyAutoGUIMouseControls") 41 | _try_load("pydirectinput", "controls.mouse.pydirectinput_mouse", "PyDirectInputMouseControls") 42 | _try_load("pynput", "controls.mouse.pynput_mouse", "PynputMouseControls") 43 | 44 | 45 | __all__ = [ 46 | "BaseMouseControls", 47 | "get_mouse_controls", 48 | "list_mouse_controls", 49 | ] 50 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from typing import Dict, Any, Optional, Tuple 3 | 4 | 5 | # Config for capture region 6 | @dataclass 7 | class CaptureRegion: 8 | left: int = 0 9 | top: int = 0 10 | width: int = 1920 11 | height: int = 1080 12 | 13 | def to_dict(self) -> Dict[str, int]: 14 | return { 15 | "left": self.left, 16 | "top": self.top, 17 | "width": self.width, 18 | "height": self.height, 19 | } 20 | 21 | def as_tuple(self) -> Tuple[int, int, int, int]: 22 | return (self.left, self.top, self.width, self.height) 23 | 24 | 25 | # Config for OBS Virtual Camera grabber 26 | @dataclass 27 | class OBSConfig: 28 | device_index: int = -1 29 | device_name: str = "OBS Virtual Camera" 30 | 31 | 32 | # Main application configuration 33 | @dataclass 34 | class AppConfig: 35 | window_title: str = "" 36 | activation_hotkey: int = 58 37 | show_preview: bool = True 38 | preview_size: Tuple[int, int] = (1280, 720) 39 | exit_on_error: bool = True 40 | 41 | grabber_type: str = "mss" 42 | grabber_options: Dict[str, Any] = field(default_factory=dict) 43 | 44 | capture_region: CaptureRegion = field(default_factory=CaptureRegion) 45 | border_offsets: Tuple[int, int, int, int] = (0, 0, 0, 0) 46 | 47 | obs: Optional[OBSConfig] = None 48 | 49 | 50 | def round_to_multiple(number: int, multiple: int) -> int: 51 | return multiple * round(number / multiple) 52 | 53 | 54 | def adjust_region_to_multiple(region: CaptureRegion, multiple: int = 32) -> CaptureRegion: 55 | return CaptureRegion( 56 | left=region.left, 57 | top=region.top, 58 | width=round_to_multiple(region.width, multiple), 59 | height=round_to_multiple(region.height, multiple), 60 | ) 61 | -------------------------------------------------------------------------------- /.idea/inspectionProfiles/Project_Default.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 41 | -------------------------------------------------------------------------------- /grabbers/win32_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Tuple 2 | 3 | import numpy as np 4 | 5 | from .base import BaseGrabber 6 | 7 | 8 | class Win32Grabber(BaseGrabber): 9 | 10 | _type = "win32" 11 | 12 | def _capture(self, region: Optional[Tuple[int, int, int, int]] = None) -> np.ndarray: 13 | import win32gui 14 | import win32ui 15 | import win32con 16 | import win32api 17 | 18 | hwin = win32gui.GetDesktopWindow() 19 | 20 | if region: 21 | left, top, x2, y2 = region 22 | width = x2 - left 23 | height = y2 - top 24 | else: 25 | width = win32api.GetSystemMetrics(win32con.SM_CXVIRTUALSCREEN) 26 | height = win32api.GetSystemMetrics(win32con.SM_CYVIRTUALSCREEN) 27 | left = win32api.GetSystemMetrics(win32con.SM_XVIRTUALSCREEN) 28 | top = win32api.GetSystemMetrics(win32con.SM_YVIRTUALSCREEN) 29 | 30 | hwindc = win32gui.GetWindowDC(hwin) 31 | srcdc = win32ui.CreateDCFromHandle(hwindc) 32 | memdc = srcdc.CreateCompatibleDC() 33 | bmp = win32ui.CreateBitmap() 34 | bmp.CreateCompatibleBitmap(srcdc, width, height) 35 | memdc.SelectObject(bmp) 36 | memdc.BitBlt((0, 0), (width, height), srcdc, (left, top), win32con.SRCCOPY) 37 | 38 | signed_ints_array = bmp.GetBitmapBits(True) 39 | img = np.frombuffer(signed_ints_array, dtype=np.uint8) 40 | img.shape = (height, width, 4) 41 | 42 | srcdc.DeleteDC() 43 | memdc.DeleteDC() 44 | win32gui.ReleaseDC(hwin, hwindc) 45 | win32gui.DeleteObject(bmp.GetHandle()) 46 | 47 | return img[:, :, :3] 48 | 49 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 50 | region = ( 51 | grab_area["left"], 52 | grab_area["top"], 53 | grab_area["left"] + grab_area["width"], 54 | grab_area["top"] + grab_area["height"], 55 | ) 56 | return self._capture(region) 57 | -------------------------------------------------------------------------------- /utils/windmouse.py: -------------------------------------------------------------------------------- 1 | """ 2 | WindMouse algorithm for human-like mouse movement. 3 | Original author: Benjamin J. Land 4 | License: GPLv3 5 | """ 6 | from typing import Callable, Tuple 7 | 8 | import numpy as np 9 | 10 | _SQRT3 = np.sqrt(3) 11 | _SQRT5 = np.sqrt(5) 12 | 13 | 14 | def wind_mouse( 15 | start_x: float, 16 | start_y: float, 17 | dest_x: float, 18 | dest_y: float, 19 | gravity: float = 9.0, 20 | wind: float = 3.0, 21 | max_step: float = 15.0, 22 | target_area: float = 12.0, 23 | move_callback: Callable[[int, int], None] = lambda x, y: None, 24 | ) -> Tuple[int, int]: 25 | current_x, current_y = start_x, start_y 26 | v_x = v_y = w_x = w_y = 0.0 27 | m_0 = max_step 28 | 29 | while True: 30 | dist = np.hypot(dest_x - start_x, dest_y - start_y) 31 | if dist < 1: 32 | break 33 | 34 | w_mag = min(wind, dist) 35 | 36 | if dist >= target_area: 37 | w_x = w_x / _SQRT3 + (2 * np.random.random() - 1) * w_mag / _SQRT5 38 | w_y = w_y / _SQRT3 + (2 * np.random.random() - 1) * w_mag / _SQRT5 39 | else: 40 | w_x /= _SQRT3 41 | w_y /= _SQRT3 42 | if m_0 < 3: 43 | m_0 = np.random.random() * 3 + 3 44 | else: 45 | m_0 /= _SQRT5 46 | 47 | v_x += w_x + gravity * (dest_x - start_x) / dist 48 | v_y += w_y + gravity * (dest_y - start_y) / dist 49 | 50 | v_mag = np.hypot(v_x, v_y) 51 | if v_mag > m_0: 52 | v_clip = m_0 / 2 + np.random.random() * m_0 / 2 53 | v_x = (v_x / v_mag) * v_clip 54 | v_y = (v_y / v_mag) * v_clip 55 | 56 | start_x += v_x 57 | start_y += v_y 58 | 59 | move_x = int(np.round(start_x)) 60 | move_y = int(np.round(start_y)) 61 | 62 | if current_x != move_x or current_y != move_y: 63 | current_x, current_y = move_x, move_y 64 | move_callback(current_x, current_y) 65 | 66 | return (current_x, current_y) 67 | -------------------------------------------------------------------------------- /grabbers/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Type, Any 2 | 3 | from .base import BaseGrabber 4 | 5 | _GRABBERS: Dict[str, Type[BaseGrabber]] = {} 6 | _LOAD_ERRORS: Dict[str, str] = {} 7 | 8 | 9 | def register_grabber(name: str): 10 | def decorator(cls: Type[BaseGrabber]): 11 | _GRABBERS[name] = cls 12 | return cls 13 | return decorator 14 | 15 | 16 | def get_grabber(name: str, **kwargs: Any) -> BaseGrabber: 17 | _ensure_grabbers_loaded() 18 | 19 | if name not in _GRABBERS: 20 | available = list(_GRABBERS.keys()) 21 | error_msg = f"Unknown grabber: {name}. Available: {available}" 22 | if name in _LOAD_ERRORS: 23 | error_msg += f"\nLoad error for '{name}': {_LOAD_ERRORS[name]}" 24 | raise ValueError(error_msg) 25 | 26 | grabber = _GRABBERS[name]() 27 | grabber.initialize(**kwargs) 28 | return grabber 29 | 30 | 31 | def list_grabbers() -> list: 32 | _ensure_grabbers_loaded() 33 | return list(_GRABBERS.keys()) 34 | 35 | 36 | def _try_load(name: str, module: str, class_name: str) -> None: 37 | try: 38 | mod = __import__(module, fromlist=[class_name]) 39 | _GRABBERS[name] = getattr(mod, class_name) 40 | except Exception as e: 41 | _LOAD_ERRORS[name] = f"{type(e).__name__}: {e}" 42 | 43 | 44 | def _ensure_grabbers_loaded(): 45 | if _GRABBERS: 46 | return 47 | 48 | _try_load("mss", "grabbers.mss_grabber", "MSSGrabber") 49 | _try_load("dxcam", "grabbers.dxcam_grabber", "DXCamGrabber") 50 | _try_load("dxcam_capture", "grabbers.dxcam_capture_grabber", "DXCamCaptureGrabber") 51 | _try_load("win32", "grabbers.win32_grabber", "Win32Grabber") 52 | _try_load("obs_vc", "grabbers.obs_vc_grabber", "OBSVirtualCameraGrabber") 53 | _try_load("d3dshot", "grabbers.d3dshot_grabber", "D3DShotGrabber") 54 | _try_load("screengear", "grabbers.screengear_grabber", "ScreenGearGrabber") 55 | 56 | 57 | __all__ = [ 58 | "BaseGrabber", 59 | "get_grabber", 60 | "list_grabbers", 61 | "register_grabber", 62 | ] 63 | -------------------------------------------------------------------------------- /grabbers/obs_vc_grabber.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, Optional, Any 2 | 3 | import cv2 4 | import numpy as np 5 | 6 | from .base import BaseGrabber 7 | from exceptions import DeviceNotFoundError 8 | 9 | 10 | class OBSVirtualCameraGrabber(BaseGrabber): 11 | 12 | _type = "obs_vc" 13 | 14 | def __init__(self): 15 | self._device: Optional[cv2.VideoCapture] = None 16 | self._size_configured = False 17 | 18 | def initialize( 19 | self, 20 | device_index: int = -1, 21 | device_name: str = "OBS Virtual Camera", 22 | **kwargs: Any, 23 | ) -> None: 24 | if device_index >= 0: 25 | self._device = cv2.VideoCapture(device_index) 26 | else: 27 | from pygrabber.dshow_graph import FilterGraph 28 | graph = FilterGraph() 29 | devices = graph.get_input_devices() 30 | 31 | try: 32 | idx = devices.index(device_name) 33 | except ValueError: 34 | raise DeviceNotFoundError( 35 | f'Device "{device_name}" not found. Available: {devices}' 36 | ) 37 | 38 | self._device = cv2.VideoCapture(idx) 39 | 40 | def _configure_size(self, width: int, height: int) -> None: 41 | if self._device is not None: 42 | self._device.set(cv2.CAP_PROP_FRAME_WIDTH, width) 43 | self._device.set(cv2.CAP_PROP_FRAME_HEIGHT, height) 44 | self._size_configured = True 45 | 46 | def get_image(self, grab_area: Dict[str, int]) -> Optional[np.ndarray]: 47 | if self._device is None: 48 | self.initialize() 49 | 50 | if not self._size_configured: 51 | self._configure_size(grab_area["width"], grab_area["height"]) 52 | 53 | ret, frame = self._device.read() 54 | if not ret or frame is None: 55 | return None 56 | 57 | return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 58 | 59 | def cleanup(self) -> None: 60 | if self._device is not None: 61 | self._device.release() 62 | self._device = None 63 | self._size_configured = False 64 | -------------------------------------------------------------------------------- /utils/cv.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, List 2 | 3 | 4 | def combine_bboxes(box1: Tuple[int, ...], box2: Tuple[int, ...]) -> Tuple[int, int, int, int]: 5 | box1_br = (box1[0] + box1[2], box1[1] + box1[3]) 6 | box2_br = (box2[0] + box2[2], box2[1] + box2[3]) 7 | 8 | x = min(box1[0], box2[0]) 9 | y = min(box1[1], box2[1]) 10 | w = max(box1_br[0], box2_br[0]) - x 11 | h = max(box1_br[1], box2_br[1]) - y 12 | 13 | return (x, y, w, h) 14 | 15 | 16 | def xywh_to_xyxy(rect: Tuple[int, ...]) -> Tuple[int, int, int, int]: 17 | return (rect[0], rect[1], rect[0] + rect[2], rect[1] + rect[3]) 18 | 19 | 20 | def xyxy_to_xywh(rect: Tuple[int, ...]) -> Tuple[int, int, int, int]: 21 | return (rect[0], rect[1], rect[2] - rect[0], rect[3] - rect[1]) 22 | 23 | 24 | def calc_iou(box_a: Tuple[int, ...], box_b: Tuple[int, ...]) -> float: 25 | x_a = max(box_a[0], box_b[0]) 26 | y_a = max(box_a[1], box_b[1]) 27 | x_b = min(box_a[2], box_b[2]) 28 | y_b = min(box_a[3], box_b[3]) 29 | 30 | inter_area = max(0, x_b - x_a) * max(0, y_b - y_a) 31 | if inter_area == 0: 32 | return 0.0 33 | 34 | box_a_area = abs((box_a[2] - box_a[0]) * (box_a[3] - box_a[1])) 35 | box_b_area = abs((box_b[2] - box_b[0]) * (box_b[3] - box_b[1])) 36 | 37 | return inter_area / float(box_a_area + box_b_area - inter_area) 38 | 39 | 40 | def boxes_intersect(box1: Tuple[int, ...], box2: Tuple[int, ...]) -> bool: 41 | return calc_iou(xywh_to_xyxy(box1), xywh_to_xyxy(box2)) > 0 42 | 43 | 44 | def merge_overlapping_boxes(boxes: List[Tuple[int, ...]]) -> List[Tuple[int, int, int, int]]: 45 | result = list(boxes) 46 | changed = True 47 | 48 | while changed: 49 | changed = False 50 | for i, box_i in enumerate(result): 51 | for j, box_j in enumerate(result): 52 | if i >= j: 53 | continue 54 | if boxes_intersect(box_i, box_j): 55 | merged = combine_bboxes(box_i, box_j) 56 | result[i] = merged 57 | result.pop(j) 58 | changed = True 59 | break 60 | if changed: 61 | break 62 | 63 | return result 64 | 65 | 66 | def point_offset(src: Tuple[int, int], dst: Tuple[int, int]) -> Tuple[int, int]: 67 | return (dst[0] - src[0], dst[1] - src[1]) 68 | 69 | 70 | def round_to_multiple(value: int, multiple: int) -> int: 71 | return multiple * round(value / multiple) 72 | -------------------------------------------------------------------------------- /controls/mouse/win32_mouse.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | from typing import Tuple 3 | 4 | from .base import BaseMouseControls 5 | 6 | 7 | class Win32MouseControls(BaseMouseControls): 8 | 9 | _type = "win32" 10 | 11 | MOUSEEVENTF_MOVE = 0x0001 12 | MOUSEEVENTF_LEFTDOWN = 0x0002 13 | MOUSEEVENTF_LEFTUP = 0x0004 14 | MOUSEEVENTF_RIGHTDOWN = 0x0008 15 | MOUSEEVENTF_RIGHTUP = 0x0010 16 | MOUSEEVENTF_MIDDLEDOWN = 0x0020 17 | MOUSEEVENTF_MIDDLEUP = 0x0040 18 | MOUSEEVENTF_WHEEL = 0x0800 19 | MOUSEEVENTF_ABSOLUTE = 0x8000 20 | SM_CXSCREEN = 0 21 | SM_CYSCREEN = 1 22 | 23 | def _do_event(self, flags: int, x_pos: int, y_pos: int, data: int = 0) -> int: 24 | x_calc = int(65536 * x_pos / ctypes.windll.user32.GetSystemMetrics(self.SM_CXSCREEN) + 1) 25 | y_calc = int(65536 * y_pos / ctypes.windll.user32.GetSystemMetrics(self.SM_CYSCREEN) + 1) 26 | return ctypes.windll.user32.mouse_event(flags, x_calc, y_calc, data, 0) 27 | 28 | def _get_button_flags(self, button: str, up: bool = False) -> int: 29 | flags = 0 30 | if "left" in button: 31 | flags |= self.MOUSEEVENTF_LEFTDOWN 32 | if "right" in button: 33 | flags |= self.MOUSEEVENTF_RIGHTDOWN 34 | if "middle" in button: 35 | flags |= self.MOUSEEVENTF_MIDDLEDOWN 36 | if up: 37 | flags <<= 1 38 | return flags 39 | 40 | def move(self, x: int, y: int) -> None: 41 | import win32api 42 | old_x, old_y = win32api.GetCursorPos() 43 | x = x if x != -1 else old_x 44 | y = y if y != -1 else old_y 45 | self._do_event(self.MOUSEEVENTF_MOVE | self.MOUSEEVENTF_ABSOLUTE, x, y) 46 | 47 | def move_relative(self, dx: int, dy: int) -> None: 48 | import win32api 49 | import win32con 50 | win32api.mouse_event(win32con.MOUSEEVENTF_MOVE, dx, dy, 0, 0) 51 | 52 | def click(self, button: str = "left") -> None: 53 | down = self._get_button_flags(button, up=False) 54 | up = self._get_button_flags(button, up=True) 55 | self._do_event(down | up, 0, 0) 56 | 57 | def get_position(self) -> Tuple[int, int]: 58 | import win32api 59 | return win32api.GetCursorPos() 60 | 61 | def press(self, button: str = "left") -> None: 62 | self._do_event(self._get_button_flags(button, up=False), 0, 0) 63 | 64 | def release(self, button: str = "left") -> None: 65 | self._do_event(self._get_button_flags(button, up=True), 0, 0) 66 | -------------------------------------------------------------------------------- /streaming/client.py: -------------------------------------------------------------------------------- 1 | import socket 2 | import time 3 | from typing import Tuple 4 | 5 | import cv2 6 | import numpy as np 7 | import lz4.frame 8 | 9 | 10 | def recv_all(conn: socket.socket, length: int) -> bytes: 11 | buf = b"" 12 | while len(buf) < length: 13 | data = conn.recv(length - len(buf)) 14 | if not data: 15 | return data 16 | buf += data 17 | return buf 18 | 19 | 20 | class StreamClient: 21 | 22 | def __init__( 23 | self, 24 | host: str = "127.0.0.1", 25 | port: int = 4000, 26 | resolution: Tuple[int, int] = (1280, 720), 27 | ): 28 | self.host = host 29 | self.port = port 30 | self.width, self.height = resolution 31 | self._socket: socket.socket = None 32 | 33 | def connect(self) -> None: 34 | self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 35 | self._socket.connect((self.host, self.port)) 36 | 37 | def disconnect(self) -> None: 38 | if self._socket: 39 | self._socket.close() 40 | self._socket = None 41 | 42 | def receive_frame(self) -> np.ndarray: 43 | size_len = int.from_bytes(self._socket.recv(1), byteorder="big") 44 | size = int.from_bytes(recv_all(self._socket, size_len), byteorder="big") 45 | compressed = recv_all(self._socket, size) 46 | pixels = lz4.frame.decompress(compressed) 47 | 48 | frame = np.frombuffer(pixels, dtype=np.uint8) 49 | frame = frame.reshape(self.height, self.width, 3) 50 | return cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) 51 | 52 | def __enter__(self): 53 | self.connect() 54 | return self 55 | 56 | def __exit__(self, exc_type, exc_val, exc_tb): 57 | self.disconnect() 58 | return False 59 | 60 | 61 | def run_viewer(host: str = "127.0.0.1", port: int = 4000) -> None: 62 | prev_time = time.time() 63 | font = cv2.FONT_HERSHEY_SIMPLEX 64 | 65 | with StreamClient(host, port) as client: 66 | while True: 67 | frame = client.receive_frame() 68 | 69 | now = time.time() 70 | fps = 1 / (now - prev_time) if now != prev_time else 0 71 | prev_time = now 72 | 73 | cv2.putText(frame, f"{int(fps)}", (7, 40), font, 1, (100, 255, 0), 3) 74 | cv2.imshow("Stream Viewer", frame) 75 | 76 | if cv2.waitKey(1) & 0xFF == ord("q"): 77 | break 78 | 79 | cv2.destroyAllWindows() 80 | 81 | 82 | if __name__ == "__main__": 83 | run_viewer() 84 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Created by .ignore support plugin (hsz.mobi) 2 | ### Python template 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | env/ 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *,cover 48 | .hypothesis/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | 58 | # Flask stuff: 59 | instance/ 60 | .webassets-cache 61 | 62 | # Scrapy stuff: 63 | .scrapy 64 | 65 | # Sphinx documentation 66 | docs/_build/ 67 | 68 | # PyBuilder 69 | target/ 70 | 71 | # IPython Notebook 72 | .ipynb_checkpoints 73 | 74 | # pyenv 75 | .python-version 76 | 77 | # celery beat schedule file 78 | celerybeat-schedule 79 | 80 | # dotenv 81 | .env 82 | 83 | # virtualenv 84 | venv/ 85 | ENV/ 86 | 87 | # Spyder project settings 88 | .spyderproject 89 | 90 | # Rope project settings 91 | .ropeproject 92 | ### VirtualEnv template 93 | # Virtualenv 94 | # http://iamzed.com/2009/05/07/a-primer-on-virtualenv/ 95 | [Bb]in 96 | [Ii]nclude 97 | [Ll]ib 98 | [Ll]ib64 99 | [Ll]ocal 100 | [Ss]cripts 101 | pyvenv.cfg 102 | .venv 103 | pip-selfcheck.json 104 | 105 | ### JetBrains template 106 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider 107 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 108 | 109 | # User-specific stuff 110 | .idea/**/workspace.xml 111 | .idea/**/tasks.xml 112 | .idea/**/usage.statistics.xml 113 | .idea/**/dictionaries 114 | .idea/**/shelf 115 | 116 | # AWS User-specific 117 | .idea/**/aws.xml 118 | 119 | # Generated files 120 | .idea/**/contentModel.xml 121 | 122 | # Sensitive or high-churn files 123 | .idea/**/dataSources/ 124 | .idea/**/dataSources.ids 125 | .idea/**/dataSources.local.xml 126 | .idea/**/sqlDataSources.xml 127 | .idea/**/dynamic.xml 128 | .idea/**/uiDesigner.xml 129 | .idea/**/dbnavigator.xml 130 | 131 | # Gradle 132 | .idea/**/gradle.xml 133 | .idea/**/libraries 134 | 135 | # Gradle and Maven with auto-import 136 | # When using Gradle or Maven with auto-import, you should exclude module files, 137 | # since they will be recreated, and may cause churn. Uncomment if using 138 | # auto-import. 139 | # .idea/artifacts 140 | # .idea/compiler.xml 141 | # .idea/jarRepositories.xml 142 | # .idea/modules.xml 143 | # .idea/*.iml 144 | # .idea/modules 145 | # *.iml 146 | # *.ipr 147 | 148 | # CMake 149 | cmake-build-*/ 150 | 151 | # Mongo Explorer plugin 152 | .idea/**/mongoSettings.xml 153 | 154 | # File-based project format 155 | *.iws 156 | 157 | # IntelliJ 158 | out/ 159 | 160 | # mpeltonen/sbt-idea plugin 161 | .idea_modules/ 162 | 163 | # JIRA plugin 164 | atlassian-ide-plugin.xml 165 | 166 | # Cursive Clojure plugin 167 | .idea/replstate.xml 168 | 169 | # SonarLint plugin 170 | .idea/sonarlint/ 171 | 172 | # Crashlytics plugin (for Android Studio and IntelliJ) 173 | com_crashlytics_export_strings.xml 174 | crashlytics.properties 175 | crashlytics-build.properties 176 | fabric.properties 177 | 178 | # Editor-based Rest Client 179 | .idea/httpRequests 180 | 181 | # Android studio 3.1+ serialized cache file 182 | .idea/caches/build_file_checksums.ser 183 | 184 | # idea folder, uncomment if you don't need it 185 | .idea -------------------------------------------------------------------------------- /example.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import multiprocessing 3 | import signal 4 | import sys 5 | import time 6 | 7 | import cv2 8 | import keyboard 9 | 10 | from config import AppConfig, CaptureRegion, OBSConfig, adjust_region_to_multiple 11 | from grabbers import get_grabber 12 | from utils.fps import FPSCounter 13 | 14 | logging.basicConfig( 15 | level=logging.INFO, 16 | format="%(asctime)s [%(levelname)s] %(message)s", 17 | ) 18 | logger = logging.getLogger(__name__) 19 | 20 | WINDOW_TITLE = "some_window" 21 | ACTIVATION_HOTKEY = 58 22 | SHOW_PREVIEW = True 23 | GRABBER_TYPE = "mss" # or "obs_vc" if you wanna use OBS Virtual Camera for a better performance (in some cases) 24 | 25 | # separate settings for OBS Virtual Camera grabber 26 | OBS_DEVICE_INDEX = -1 27 | OBS_DEVICE_NAME = "OBS Virtual Camera" 28 | 29 | 30 | def create_config() -> AppConfig: 31 | try: 32 | from utils.win32 import get_window_rect 33 | rect = get_window_rect(WINDOW_TITLE) 34 | region = CaptureRegion( 35 | left=rect[0], 36 | top=rect[1], 37 | width=rect[2], 38 | height=rect[3], 39 | ) 40 | except Exception: 41 | region = CaptureRegion() 42 | 43 | region = adjust_region_to_multiple(region, 32) 44 | 45 | config = AppConfig( 46 | window_title=WINDOW_TITLE, 47 | activation_hotkey=ACTIVATION_HOTKEY, 48 | show_preview=SHOW_PREVIEW, 49 | grabber_type=GRABBER_TYPE, 50 | capture_region=region, 51 | ) 52 | 53 | if GRABBER_TYPE == "obs_vc": 54 | config.obs = OBSConfig( 55 | device_index=OBS_DEVICE_INDEX, 56 | device_name=OBS_DEVICE_NAME, 57 | ) 58 | config.grabber_options = { 59 | "device_index": OBS_DEVICE_INDEX, 60 | "device_name": OBS_DEVICE_NAME, 61 | } 62 | 63 | return config 64 | 65 | 66 | def grab_process( 67 | queue: multiprocessing.JoinableQueue, 68 | stop_event: multiprocessing.Event, 69 | config: AppConfig, 70 | ) -> None: 71 | try: 72 | grabber = get_grabber(config.grabber_type, **config.grabber_options) 73 | except Exception as e: 74 | logger.error(f"Failed to initialize grabber: {e}") 75 | stop_event.set() 76 | return 77 | 78 | grab_area = config.capture_region.to_dict() 79 | 80 | while not stop_event.is_set(): 81 | try: 82 | img = grabber.get_image(grab_area) 83 | if img is None: 84 | continue 85 | 86 | while not queue.empty(): 87 | try: 88 | queue.get_nowait() 89 | except BaseException: 90 | break 91 | 92 | queue.put_nowait(img) 93 | queue.join() 94 | 95 | except Exception as e: 96 | logger.error(f"Capture error: {e}") 97 | stop_event.set() 98 | break 99 | 100 | grabber.cleanup() 101 | logger.info("Capture process stopped") 102 | 103 | 104 | def process_frame( 105 | queue: multiprocessing.JoinableQueue, 106 | stop_event: multiprocessing.Event, 107 | activated: multiprocessing.Event, 108 | config: AppConfig, 109 | ) -> None: 110 | fps = FPSCounter() 111 | font = cv2.FONT_HERSHEY_SIMPLEX 112 | 113 | while not stop_event.is_set(): 114 | if queue.empty(): 115 | time.sleep(0.001) 116 | continue 117 | 118 | try: 119 | img = queue.get_nowait() 120 | queue.task_done() 121 | except BaseException: 122 | continue 123 | 124 | if activated.is_set(): 125 | pass 126 | 127 | if config.show_preview: 128 | current_fps = fps() 129 | cv2.putText( 130 | img, 131 | f"{current_fps:.1f}", 132 | (20, 120), 133 | font, 134 | 1.7, 135 | (0, 255, 0), 136 | 7, 137 | cv2.LINE_AA, 138 | ) 139 | cv2.imshow("Preview", cv2.resize(img, config.preview_size)) 140 | 141 | if cv2.waitKey(1) & 0xFF == ord("q"): 142 | stop_event.set() 143 | 144 | cv2.destroyAllWindows() 145 | logger.info("Processing stopped") 146 | 147 | 148 | def main() -> int: 149 | config = create_config() 150 | 151 | stop_event = multiprocessing.Event() 152 | activated = multiprocessing.Event() 153 | queue = multiprocessing.JoinableQueue() 154 | 155 | def toggle_activation(): 156 | if activated.is_set(): 157 | activated.clear() 158 | logger.info("Deactivated") 159 | else: 160 | activated.set() 161 | logger.info("Activated") 162 | 163 | keyboard.add_hotkey(config.activation_hotkey, toggle_activation) 164 | 165 | def shutdown(signum, frame): 166 | logger.info("Shutting down...") 167 | stop_event.set() 168 | 169 | signal.signal(signal.SIGINT, shutdown) 170 | signal.signal(signal.SIGTERM, shutdown) 171 | 172 | processes = [ 173 | multiprocessing.Process( 174 | target=grab_process, 175 | args=(queue, stop_event, config), 176 | ), 177 | multiprocessing.Process( 178 | target=process_frame, 179 | args=(queue, stop_event, activated, config), 180 | ), 181 | ] 182 | 183 | for p in processes: 184 | p.start() 185 | 186 | try: 187 | while not stop_event.is_set(): 188 | time.sleep(0.1) 189 | except KeyboardInterrupt: 190 | stop_event.set() 191 | 192 | for p in processes: 193 | p.join(timeout=3) 194 | if p.is_alive(): 195 | p.terminate() 196 | 197 | logger.info("Shutdown complete") 198 | return 0 199 | 200 | 201 | if __name__ == "__main__": 202 | sys.exit(main()) 203 | --------------------------------------------------------------------------------