78 lines
2.6 KiB
Python
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)) == BLACK:
|
|
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"))
|