dark_helper/capture_monitor.py at master · Spain-AI/dark_helper · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from mss import mss
import cv2
from PIL import Image
import numpy as np
from threading import Lock

from PIL import Image
from Xlib import X
import ewmh

class CaptureMonitor:
    MODE_CROP, MODE_WINDOW = list(range(2))
    def __init__(self, bb):
        self.lock = Lock()
        self.screen_capture = mss()
        self.ewmh_capture = ewmh.EWMH()

        self.current_window = None
        self.current_monitor = {'left': None, 'top': None, 'width': None, 'height': None}

        self.mode = CaptureMonitor.MODE_CROP
        self.change_coords(bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1])

        self.max_resolution = 920
        self.min_resolution = 320

    def screen_resize(self, im, size_axis):
        if im.shape[0] > im.shape[1]:
            h = size_axis
            w = h * (im.shape[1] / im.shape[0])
        else:
            w = size_axis
            h = w * (im.shape[0] / im.shape[1])

        h = int(h)
        w = int(w)
        return cv2.resize(im, (w, h))

    def get_frame(self):
        with self.lock:
            if self.mode == CaptureMonitor.MODE_WINDOW and self.current_window is not None:
                geo = self.current_window.get_geometry()
                raw_im = self.current_window.get_image(0, 0, geo.width, geo.height, X.ZPixmap, 0xffffffff)
                im = Image.frombytes("RGB", (geo.width, geo.height), raw_im.data, "raw", "BGRX")
            else:
                im = Image.frombytes('RGB', self.size, self.screen_capture.grab(self.current_monitor).rgb)

            im = np.array(im) # RGB

            if im.shape[0] < self.min_resolution or im.shape[1] < self.min_resolution:
                return self.screen_resize(im, self.min_resolution)

            elif im.shape[0] > self.max_resolution or im.shape[1] > self.max_resolution:
                return self.screen_resize(im, self.max_resolution)
            return im


    def get_windows(self):
        windows = []
        for w in self.ewmh_capture.getClientList():
            name = w.get_wm_name().strip()
            if name == '':
                continue

            windows.append(name)

        return ['Active window'] + windows

    def set_window(self, name):
        with self.lock:
            if name == 'Active window':
                self.current_window = self.ewmh_capture.getActiveWindow()
                return

            for w in self.ewmh_capture.getClientList():
                if name == w.get_wm_name().strip():
                    self.current_window = w
                    return

    def set_mode(self, mode):
        with self.lock:
            self.mode = mode

    def change_coords(self, left, top, width, height):
        with self.lock:
            left = int(left)
            top = int(top)
            width = int(width)
            height = int(height)

            self.size = width, height
            self.current_monitor.update({'left': left, 'top': top, 'width': self.size[0], 'height': self.size[1]})