1 """A wrapper around the Generic Buffer Manager (GBM) library. 2 3 Currently implements exactly the functions required to screenshot a frame 4 buffer using DRM crtc info. 5 """ 6 from ctypes import * 7 import drm 8 from PIL import Image 9 10 GBM_BO_IMPORT_FD = 0x5503 11 GBM_BO_USE_SCANOUT = c_uint(1) 12 GBM_BO_TRANSFER_READ = c_uint(1) 13 GBM_MAX_PLANES = 4 14 15 def __gbm_fourcc_code(a, b, c, d): 16 return ord(a) | (ord(b) << 8) | (ord(c) << 16) | (ord(d) << 24) 17 18 GBM_FORMAT_ARGB8888 = __gbm_fourcc_code("A", "R", "2", "4") 19 GBM_LIBRARIES = ["libgbm.so", "libgbm.so.1"] 20 21 22 class gbm_import_fd_data(Structure): 23 _fields_ = [ 24 ("fd", c_int), 25 ("width", c_uint), 26 ("height", c_uint), 27 ("stride", c_uint), 28 ("bo_format", c_uint), 29 ] 30 31 32 class gbm_import_fd_planar_data(Structure): 33 _fields_ = [ 34 ("fds", c_int * GBM_MAX_PLANES), 35 ("width", c_uint), 36 ("height", c_uint), 37 ("bo_format", c_uint), 38 ("strides", c_uint * GBM_MAX_PLANES), 39 ("offsets", c_uint * GBM_MAX_PLANES), 40 ("format_modifiers", c_ulonglong * GBM_MAX_PLANES), 41 ] 42 43 44 class gbm_device(Structure): 45 """Opaque struct for GBM device. 46 """ 47 pass 48 49 50 class gbm_bo(Structure): 51 """Opaque struct for GBM buffer. 52 """ 53 pass 54 55 56 def loadGBM(): 57 """Load and return a handle to libgbm.so. 58 """ 59 l = None 60 61 for lib in GBM_LIBRARIES: 62 try: 63 l = cdll.LoadLibrary(lib) 64 except OSError: 65 l = None 66 if l is not None: 67 break 68 69 if l is None: 70 raise RuntimeError("Could not load GBM library.") 71 return None 72 73 l.gbm_create_device.argtypes = [c_int] 74 l.gbm_create_device.restype = POINTER(gbm_device) 75 76 l.gbm_device_destroy.argtypes = [POINTER(gbm_device)] 77 l.gbm_device_destroy.restype = None 78 79 l.gbm_bo_import.argtypes = [POINTER(gbm_device), c_uint, c_void_p, c_uint] 80 l.gbm_bo_import.restype = POINTER(gbm_bo) 81 82 l.gbm_bo_map.argtypes = [ 83 POINTER(gbm_bo), c_uint, c_uint, c_uint, c_uint, c_uint, 84 POINTER(c_uint), 85 POINTER(c_void_p), c_size_t 86 ] 87 l.gbm_bo_map.restype = c_void_p 88 89 l.gbm_bo_unmap.argtypes = [POINTER(gbm_bo), c_void_p] 90 l.gbm_bo_unmap.restype = None 91 92 return l 93 94 95 class GBMBuffer(object): 96 """A GBM buffer. 97 """ 98 99 def __init__(self, library, buffer): 100 self._l = library 101 self._buffer = buffer 102 103 def __del__(self): 104 if self._l: 105 self._l.gbm_bo_destroy(self._buffer) 106 107 @classmethod 108 def fromFD(cls, device, fd, width, height, stride, bo_format, usage): 109 """Create/import a GBM Buffer Object from a file descriptor. 110 111 @param device: GBM device object. 112 @param fd: a file descriptor for the buffer to be imported. 113 @param width: buffer width in pixels. 114 @param height: buffer height in pixels. 115 @param stride: buffer pitch; number of pixels between sequential rows. 116 @param bo_format: pixel format fourcc code, e.g. GBM_FORMAT_ARGB8888. 117 @param usage: buffer usage flags, e.g. GBM_BO_USE_SCANOUT. 118 """ 119 bo_data = gbm_import_fd_data() 120 bo_data.fd = fd 121 bo_data.width = width 122 bo_data.height = height 123 bo_data.stride = stride 124 bo_data.bo_format = bo_format 125 buffer = device._l.gbm_bo_import(device._device, GBM_BO_IMPORT_FD, 126 byref(bo_data), usage) 127 if buffer is None: 128 raise RuntimeError("gbm_bo_import() returned NULL") 129 130 self = cls(device._l, buffer) 131 return self 132 133 def map(self, x, y, width, height, flags, plane): 134 """Map buffer data into this user-space. 135 Returns (address, stride_bytes): void* start address for pixel array, 136 number of BYTES between sequental rows of pixels. 137 @param flags: The union of the GBM_BO_TRANSFER_* flags for this buffer. 138 """ 139 self._map_p = c_void_p(0) 140 stride_out = c_uint(0) 141 if width == 0 or height == 0: 142 raise RuntimeError("Map width and/or height is 0") 143 map_p = self._l.gbm_bo_map(self._buffer, x, y, width, height, flags, 144 byref(stride_out), byref(self._map_p), plane) 145 if stride_out is 0: 146 raise RuntimeError("gbm_bo_map() stride is 0") 147 if map_p is 0: 148 raise RuntimeError("gbm_bo_map() returned NULL") 149 return map_p, stride_out 150 151 def unmap(self, map_data_p): 152 self._l.gbm_bo_unmap(self._buffer, map_data_p) 153 154 155 class GBMDevice(object): 156 """A GBM device. 157 """ 158 159 def __init__(self, library, handle): 160 self._l = library 161 self._handle = handle 162 self._device = library.gbm_create_device(self._handle) 163 164 def __del__(self): 165 if self._l: 166 self._l.gbm_device_destroy(self._device) 167 168 @classmethod 169 def fromHandle(cls, handle): 170 """Create a device object from an open file descriptor. 171 """ 172 self = cls(loadGBM(), handle) 173 return self 174 175 176 def _bgrx24(i): 177 b = i & 255 178 g = (i >> 8) & 255 179 r = (i >> 16) & 255 180 return r, g, b 181 182 183 def crtcScreenshot(crtc_id=None): 184 """Take a screenshot, returning an image object. 185 186 @param crtc_id: The CRTC to screenshot. 187 None for first found CRTC with mode set 188 or "internal" for crtc connected to internal LCD 189 or "external" for crtc connected to external display 190 or "usb" "evdi" or "udl" for crtc with valid mode on evdi 191 or udl display 192 or DRM integer crtc_id 193 """ 194 crtc = drm.getCrtc(crtc_id) 195 if crtc is not None: 196 device = GBMDevice.fromHandle(drm._drm._fd) 197 framebuffer = crtc.fb() 198 # TODO(djmk): The buffer format is hardcoded to ARGB8888, we should fix 199 # this to query for the frambuffer's format instead. 200 format_hardcode = GBM_FORMAT_ARGB8888 201 202 bo = GBMBuffer.fromFD(device, 203 framebuffer.getFD(), framebuffer.width, 204 framebuffer.height, framebuffer.pitch, 205 format_hardcode, GBM_BO_USE_SCANOUT) 206 map_void_p, stride_bytes = bo.map(0, 0, framebuffer.width, 207 framebuffer.height, 208 GBM_BO_TRANSFER_READ, 0) 209 map_bytes = stride_bytes.value * framebuffer.height 210 211 # Create a Python Buffer object which references (but does not own) the 212 # memory. 213 buffer_from_memory = pythonapi.PyBuffer_FromMemory 214 buffer_from_memory.restype = py_object 215 buffer_from_memory.argtypes = [c_void_p, c_ssize_t] 216 map_buffer = buffer_from_memory(map_void_p, map_bytes) 217 218 # Make a copy of the bytes. Doing this is faster than the conversion, 219 # and is more likely to capture a consistent snapshot of the framebuffer 220 # contents, as a process may be writing to it. 221 buffer_bytes = bytes(map_buffer) 222 223 # Load the image, converting from the BGRX format to a PIL Image in RGB 224 # form. As the conversion is implemented by PIL as C code, this 225 # conversion is much faster than calling _bgrx24(). 226 image = Image.fromstring( 227 'RGB', (framebuffer.width, framebuffer.height), buffer_bytes, 228 'raw', 'BGRX', stride_bytes.value, 1) 229 bo.unmap(bo._map_p) 230 del bo 231 return image 232 233 raise RuntimeError( 234 "Unable to take screenshot. There may not be anything on the screen.") 235 236 237 def crtcGetPixel(x, y, crtc_id=None): 238 """Get a pixel from the specified screen, returning a rgb tuple. 239 240 @param x: The x-coordinate of the desired pixel. 241 @param y: The y-coordinate of the desired pixel. 242 @param crtc_id: The CRTC to get the pixel from. 243 None for first found CRTC with mode set 244 or "internal" for crtc connected to internal LCD 245 or "external" for crtc connected to external display 246 or "usb" "evdi" or "udl" for crtc with valid mode on evdi 247 or udl display 248 or DRM integer crtc_id 249 """ 250 crtc = drm.getCrtc(crtc_id) 251 if crtc is None: 252 raise RuntimeError( 253 "Unable to get pixel. There may not be anything on the screen.") 254 device = GBMDevice.fromHandle(drm._drm._fd) 255 framebuffer = crtc.fb() 256 bo = GBMBuffer.fromFD(device, 257 framebuffer.getFD(), framebuffer.width, 258 framebuffer.height, framebuffer.pitch, 259 GBM_FORMAT_ARGB8888, GBM_BO_USE_SCANOUT) 260 map_void_p, _ = bo.map(int(x), int(y), 1, 1, GBM_BO_TRANSFER_READ, 0) 261 map_int_p = cast(map_void_p, POINTER(c_int)) 262 pixel = _bgrx24(map_int_p[0]) 263 bo.unmap(bo._map_p) 264 return pixel 265
