152 lines
4.1 KiB
Python
152 lines
4.1 KiB
Python
from PIL import Image
|
|
from pyzbar.pyzbar import decode
|
|
import itertools
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
from multiprocessing import Pool
|
|
from math import perm
|
|
import sys
|
|
|
|
BOX_SIZE = 2
|
|
TILE = 78
|
|
# 7x7 finder, upscaled to 14x14 (black=1, white=0)
|
|
TEMPLATE = np.kron(np.array([
|
|
[1,1,1,1,1,1,1],
|
|
[1,0,0,0,0,0,1],
|
|
[1,0,1,1,1,0,1],
|
|
[1,0,1,1,1,0,1],
|
|
[1,0,1,1,1,0,1],
|
|
[1,0,0,0,0,0,1],
|
|
[1,1,1,1,1,1,1],
|
|
], dtype=np.uint8), np.ones((BOX_SIZE, BOX_SIZE), np.uint8))
|
|
H, W = TEMPLATE.shape # 14x14
|
|
|
|
|
|
def decodeQRCode(image : Image.Image) -> str | None:
|
|
# Decode the QR code
|
|
decoded_objects = decode(image)
|
|
|
|
# Check if any QR code was found
|
|
if decoded_objects:
|
|
# Iterate through all the detected QR codes and print their data
|
|
for obj in decoded_objects:
|
|
return obj.data.decode("utf-8")
|
|
else:
|
|
return None
|
|
|
|
|
|
def cut_pieces(image):
|
|
# Get the size of the image
|
|
width, height = image.size
|
|
|
|
if width == height:
|
|
NUM_PIECES_WIDTH = 3
|
|
else:
|
|
NUM_PIECES_WIDTH = 4
|
|
|
|
NUM_PIECES_HEIGHT = 3
|
|
|
|
pieces = []
|
|
|
|
for x, y in [(x, y) for x in range(NUM_PIECES_WIDTH) for y in range(NUM_PIECES_HEIGHT)]:
|
|
pieces.append(image.crop((x * TILE, y * TILE, x * TILE + TILE, y * TILE + TILE)))
|
|
|
|
return pieces
|
|
|
|
|
|
def to_bw01(img):
|
|
# Force 1-bit, then map black->1, white->0
|
|
a = np.array(img.convert("1")) # {0,255}
|
|
return (a == 0).astype(np.uint8)
|
|
|
|
|
|
def finder_corner_exact(tile_bw):
|
|
crops = {
|
|
'TL': tile_bw[0:H, 0:W ],
|
|
'TR': tile_bw[0:H, TILE-W:TILE ],
|
|
'BL': tile_bw[TILE-H:TILE, 0:W ],
|
|
}
|
|
for corner, crop in crops.items():
|
|
if np.array_equal(crop, TEMPLATE):
|
|
return corner
|
|
return None # no exact finder in this tile
|
|
|
|
|
|
def three_tiles_with_finders(tiles):
|
|
hits = []
|
|
for tile in tiles:
|
|
corner = finder_corner_exact(to_bw01(tile))
|
|
if corner is not None:
|
|
hits.append((tile, corner))
|
|
if len(hits) == 3:
|
|
break
|
|
return hits
|
|
|
|
def solve_combo(combo):
|
|
"""
|
|
Function to be executed by each process in the pool.
|
|
"""
|
|
a, b, c, d, e, f = combo
|
|
# Assuming `new_image` is a base image you can copy or an in-memory image object.
|
|
# You may need to pass `new_image` to the function or have it loaded by each process.
|
|
local_image = new_image.copy() # Make a copy to avoid race conditions
|
|
local_image.paste(a, (78, 0))
|
|
local_image.paste(b, (0, 78))
|
|
local_image.paste(c, (78, 78))
|
|
local_image.paste(d, (156, 78))
|
|
local_image.paste(e, (78, 156))
|
|
local_image.paste(f, (156, 156))
|
|
|
|
output = decodeQRCode(local_image)
|
|
if output:
|
|
return output
|
|
return None
|
|
|
|
|
|
def main(img_path="./output/secret_message.png", verbose=True):
|
|
img = Image.open(img_path)
|
|
global new_image
|
|
new_image = Image.new("RGB", (img.height, img.height))
|
|
|
|
|
|
pieces = cut_pieces(img)
|
|
corners = three_tiles_with_finders(pieces)
|
|
|
|
for tile, corner in corners:
|
|
if corner == 'TL':
|
|
new_image.paste(tile, (0,0))
|
|
elif corner == 'TR':
|
|
new_image.paste(tile, (156,0))
|
|
elif corner == 'BL':
|
|
new_image.paste(tile, (0,156))
|
|
else:
|
|
sys.exit()
|
|
|
|
|
|
pieces = [piece for piece in pieces if piece not in [tile for tile, _ in corners]]
|
|
possible_combos = itertools.permutations(pieces, 6)
|
|
|
|
result = None
|
|
|
|
# Use `multiprocessing.Pool` with the number of available cores
|
|
with Pool() as pool:
|
|
# `imap_unordered` returns results as they are completed.
|
|
# This is a lazy iterator, so it doesn't create all combos at once.
|
|
# `tqdm` can still be used to show progress.
|
|
for result in tqdm(pool.imap_unordered(solve_combo, possible_combos), total=perm(len(pieces), 6), disable=(not verbose)):
|
|
if result:
|
|
if verbose:
|
|
print(f"QR Code decoded: {result}")
|
|
pool.terminate() # Stop all worker processes
|
|
break
|
|
|
|
return result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) == 2:
|
|
main(sys.argv[1])
|
|
else:
|
|
main()
|
|
|