├── 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 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
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 |
4 |
5 |
16 |
17 |
18 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
40 |
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 |
--------------------------------------------------------------------------------