mower-ng/mower/utils/qrcode.py

78 lines
2.6 KiB
Python

import json
from typing import Dict, List, Optional
from zlib import compress, decompress
from base45 import b45decode, b45encode
from PIL import Image, ImageChops, ImageDraw
from pyzbar import pyzbar
from qrcode.constants import ERROR_CORRECT_L
from qrcode.main import QRCode
QRCODE_SIZE = 215
GAP_SIZE = 16
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
TOP = 40
BOTTOM = 995
LEFT = 40
def encode(data: str, n: int = 16, theme: str = "light") -> List[Image.Image]:
data = b45encode(compress(data.encode("utf-8"), level=9))
length = len(data)
split: List[bytes] = []
for i in range(n):
start = length // n * i
end = length if i == n - 1 else length // n * (i + 1)
split.append(data[start:end])
result: List[Image.Image] = []
qr = QRCode(error_correction=ERROR_CORRECT_L)
fg, bg = (BLACK, WHITE) if theme == "light" else (WHITE, BLACK)
for i in split:
qr.add_data(i)
img: Image.Image = qr.make_image(fill_color=fg, back_color=bg)
result.append(trim(img.get_image()))
qr.clear()
return result
def trim(img: Image.Image) -> Image.Image:
bg = Image.new(img.mode, img.size, img.getpixel((0, 0)))
diff = ImageChops.difference(img, bg)
img = img.crop(diff.getbbox())
img = img.resize((QRCODE_SIZE, QRCODE_SIZE))
return img
def export(plan: Dict, img: Image.Image, theme: str = "light") -> Image.Image:
qrcode_list = encode(json.dumps(plan), theme=theme)
for idx, i in enumerate(qrcode_list[:7]):
img.paste(i, (LEFT + idx * (GAP_SIZE + QRCODE_SIZE), TOP))
for idx, i in enumerate(qrcode_list[7:14]):
img.paste(i, (LEFT + idx * (GAP_SIZE + QRCODE_SIZE), BOTTOM))
for idx, i in enumerate(qrcode_list[14:]):
img.paste(i, (2520 + idx * (GAP_SIZE + QRCODE_SIZE), BOTTOM))
img = img.convert("RGB")
return img
def decode(img: Image.Image) -> Optional[Dict]:
img = img.convert("RGB")
if img.getpixel((0, 0))[0] < 127:
img = ImageChops.invert(img)
result = []
while len(data := pyzbar.decode(img)):
img1 = ImageDraw.Draw(img)
for d in data:
if d.quality > 1:
continue
left = d.rect.left - 2
top = d.rect.top - 2
right = left + d.rect.width + 5
bottom = top + d.rect.height + 5
scope = ((left, top), (right, bottom))
img1.rectangle(scope, fill=WHITE)
result.append(d)
result.sort(key=lambda i: (i.rect.top * 2 > img.size[1], i.rect.left))
result = b45decode(b"".join([i.data for i in result]))
return json.loads(decompress(result).decode("utf-8"))