PIL.ImageGrabの仕様がひどすぎたのでスクリーンショットを自作した

PIL.ImageGrabの仕様がひどすぎたのでスクリーンショットを自作した

はじめに

PythonのフレームワークにPILというものがあります。

これは画像を処理したりするのによく使われます。

また、PIL.ImageGrabをインポートすると、スクリーンショットができます。

が、この仕様がひどい。

何がひどいかというと……、まあ見てください。

def grab(bbox=None, include_layered_windows=False, all_screens=False, xdisplay=None):
    if xdisplay is None:
        if sys.platform == "darwin":
            fh, filepath = tempfile.mkstemp(".png")
            os.close(fh)
            subprocess.call(["screencapture", "-x", filepath])
            im = Image.open(filepath)
            im.load()
            os.unlink(filepath)
            if bbox:
                im_cropped = im.crop(bbox)
                im.close()
                return im_cropped
            return im
        elif sys.platform == "win32":
            offset, size, data = Image.core.grabscreen_win32(
                include_layered_windows, all_screens
            )
            im = Image.frombytes(
                "RGB",
                size,
                data,
                # RGB, 32-bit line padding, origin lower left corner
                "raw",
                "BGR",
                (size[0] * 3 + 3) & -4,
                -1,
            )
            if bbox:
                x0, y0 = offset
                left, top, right, bottom = bbox
                im = im.crop((left - x0, top - y0, right - x0, bottom - y0))
            return im
    # use xdisplay=None for default display on non-win32/macOS systems
    if not Image.core.HAVE_XCB:
        raise OSError("Pillow was built without XCB support")
    size, data = Image.core.grabscreen_x11(xdisplay)
    im = Image.frombytes("RGB", size, data, "raw", "BGRX", size[0] * 4, 1)
    if bbox:
        im = im.crop(bbox)
    return im

subprocessでスクリーンショットするのはまだ許せます。

が、screencaptureにはそもそも-Rという、矩形領域を指定して画像を取得するという機能があるのです。

なのにそれを使わずに、わざわざ自前のim.crop()を使って画像をクロップしているのです。

つまり、私の所望である「矩形領域選択→スクリーンショットしたデータをNumPyなりで処理」ではなく、「全画面スクリーンショット→画像を保存→画像の読み込み→クロップ」という方式をとっているのです。

どうりで速度が遅いわけだ。

NumPyで処理する方が速度が速いしさらに後々便利なので、cv2で処理できれば嬉しいな、という感じで調べましたが、残念ながらcv2にはスクリーンショット自体がない。

じゃあどうするか。作るしかないでしょう。

環境

  • Catalina
  • Python3.8.5
  • numpy==1.19.5
  • cv2==4.1.2

作っている途中の進捗状態

最初に困ったのはscreencapture

ターミナルでこのコマンドを実行すると全画面スクリーンショットができるのですが、どうやって一部の矩形領域をキャプチャーするものか考えていました。

そこで見つけたのが、このサイト

別に英語が読めないわけではないけど、-hで見るのも最近は忙しすぎて辛かったので、サイトに頼ることにした。

んで、どうやら-R <x,y,w,h>で指定するとみられるらしい、ということがわかりました。

また、ここで見つけたのは、どうやらファイルを保存してから読み出すとかいうゴミムーヴをしなくても済む-cという便利な引数があるという。嬉しい。

ともあれ実現。

import subprocess

subprocess.run(f'screencapture -c -R {x},{y},{w},{h}'.split(' '))

次に困ったのは画像サイズ。

subprocessでちゃんと動くための引数を割り当てないといけませんが、x, y, w, hの最小値、最大値とか、どうすりゃ良いのさ、となったのときに見つけたのが、Pythonのフレームワークであるscreeninfo。

こいつはスクリーンのサイズを取り出してくれる、らしい。

これをとりあえずpip installしてみて、screeninfo/enumerator/osx.pyの中身を見てみた。

import typing as T

from screeninfo.common import Monitor


def enumerate_monitors() -> T.Iterable[Monitor]:
    from AppKit import NSScreen

    screens = NSScreen.screens()

    for screen in screens:
        f = screen.frame
        if callable(f):
            f = f()

        yield Monitor(
            x=int(f.origin.x),
            y=int(f.origin.y),
            width=int(f.size.width),
            height=int(f.size.height),
        )

なるほど、AppKitを使って見てるのか。

とまあ、これを参考にして組み立てる。

import sys
if sys.platform != 'darwin': raise OSError(f'unsupported OS.<{sys.platform=}>')
import subprocess
import numpy as np
import AppKit
import cv2


def rect(
    x: int or float or str,
    y: int or float or str,
    w: int or float or str,
    h: int or float or str
):
    x, y, w, h = map(float, [x, y, w, h])

    screens = AppKit.NSScreen.screens()
    if len(screens) > 1:
        raise SystemError(
            'unsupported more than one screens.'
            f'<{len(screens)=}>'
        )

    f = screens[0].frame
    if callable(f):
        f = f()

    x_max, y_max = f.size.width, f.size.height
    w_max, h_max = x_max - x - w, y_max - y - h

    if 0 <= x < x_max and 0 <= y < y_max and 0 <= w <= w_max and 0 <= h <= h_max:
        subprocess.run(f'screencapture -c -R {x},{y},{w},{h}'.split(' '))

どうやらマルチスクリーンに対応してくれてるらしいけど、さすがに大変なのでやめた。

また、screencaptureで引数に-cを良くも悪くも与えてしまったので、スクリーンショットで取得したデータがクリップボードに入っている。

そのデータをbytes()してnumpy.frombuffer()してcv2.imdecode()する必要があった。

でも保存するよりは速いので書く。

プログラム

で、できたのが下のプログラムです。これををscreencapture/__init__.pyに保存しました。

import sys
if sys.platform != 'darwin': raise OSError(f'unsupported OS.<{sys.platform=}>')
import subprocess
import numpy as np
import AppKit
import cv2


def rect(
    x: int or float or str,
    y: int or float or str,
    w: int or float or str,
    h: int or float or str
):
    x, y, w, h = map(float, [x, y, w, h])

    screens = AppKit.NSScreen.screens()
    if len(screens) > 1:
        raise SystemError(
            'unsupported more than one screens.'
            f'<{len(screens)=}>'
        )

    f = screens[0].frame
    if callable(f):
        f = f()

    x_max, y_max = f.size.width, f.size.height
    w_max, h_max = x_max - x - w, y_max - y - h

    if 0 <= x < x_max and 0 <= y < y_max and 0 <= w <= w_max and 0 <= h <= h_max:
        subprocess.run(f'screencapture -c -R {x},{y},{w},{h}'.split(' '))
        board = AppKit.NSPasteboard.generalPasteboard()

        return cv2.imdecode(np.frombuffer(bytes(board.dataForType_(board.types()[0])), np.uint8), -1)

    else:
        raise ValueError(
            f'arguments must be 0<=x<{int(x_max)}, 0<=y<{int(y_max)}, 0<=w<={int(w_max)}, 0<=h<={int(h_max)}.'
            f'<{(x, y, w, h)=}>'
        )


if __name__ == '__main__':

    x, y, w, h = 400, 200, 300, 180
    rect(x, y, w, h)

で、実行ファイルはこれ。

import screencapture as scap
import cv2

if __name__ == '__main__':

    x, y, w, h = 400, 200, 300, 180
    cv2.imwrite('sample.png', scap.rect(x, y, w, h))

これを実行すると座標 \( (x,y)=(400,200) \)から\( 300 \times 180 \)の画像をNumPyの形式で取得して、それをcv2で保存することができます。

問題点

まず、私自身がmacユーザーなので、sys.platform=='darwin'以外定義していないということです。

Windowsは自分たちで作ってください。

また、スクリーンが複数ある場合には対応していません。

デュアルモニターとかトリプルモニターとかそういうのにも対応していません。

さらに、これはcv2の問題ですが、画像がBGRAで出力されます。ま、PILがcv2に負ければ解決する問題だけどね。

さいごに

高速化のために結構手間をかけましたが、AppKitがどういう構造になっているのかということを知ることができたし、何しろ他人が書いたフレームワークのソースコードを読めるようになっていたことが嬉しかったです。

3年前からは想像できなかった進捗だね。

また、面白いものを見つけ次第、書きます。