uiautomator2刘海屏截图代码

notch_screenshot.py

# coding: utf-8
#

from PIL import Image, ImageDraw

class Notch(object):
    def __init__(self, d):
        self.d = d
        self.brand = self.d.device_info['brand']
        self.model = self.d.device_info['model']
        self.color = (255, 0, 0) # RGB
        print(self.brand, self.model)
        if self.brand == 'Xiaomi' and self.model == 'MI 8':
            print("Good")
            self.notch = True
            self.notch_height = 89
            self.notch_width = 560
            self.notch_diameter = 200 # 底部的圆直径
        else:
            self.notch = False

    def __call__(self):
        if not self.notch:
            return self.d.screenshot()
        import time
        start = time.time()
        rdict = {1: Image.ROTATE_90, 2: Image.ROTATE_180, 3: Image.ROTATE_270}
        im = d.screenshot()
        rotation = self.d.info.get('displayRotation')
        if rotation:
            im = im.transpose(rdict[4-rotation])
        draw = ImageDraw.Draw(im)
        width = self.notch_width
        height = self.notch_height

        # 画刘海
        lx = (im.size[0] - width) / 2
        rx = lx + width
        draw.line((lx, height, rx, height), fill=self.color, width=5)
        draw.line((lx-height/2, 0, lx, height), fill=self.color, width=5)
        draw.line((rx+height/2, 0, rx, height), fill=self.color, width=5)

        # 画四周的圆弧
        dia = self.notch_diameter # 直径
        draw.arc([0, 0, dia, dia], 180, 270, fill=self.color)
        w, h = im.size
        draw.arc([w-dia, 0, w, dia], 270, 0, fill=self.color)
        draw.arc([0, h-dia, dia, h], 90, 180, fill=self.color)
        draw.arc([w-dia, h-dia, w, h], 0, 90, fill=self.color)

        del draw
        if rotation:
            im = im.transpose(rdict[rotation])
        print("time used:", time.time() - start)
        return im


if __name__ == '__main__':
    import uiautomator2 as u2
    u2.plugin_register('notch', Notch)
    d = u2.connect()
    d.screenshot("h1.jpg")
    d.ext_notch().save("haha.jpg")