Misc - Puzzle 2 - OpenCV Solver

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import cv2
import numpy as np
from math import sqrt
from multiprocessing import Pool, cpu_count

FOND_PATH    = 'fond.jpeg'
TILES_DIR    = 'tiles'
BORDER_PX    = 8
SCALE_COARSE = 0.25
WINDOW_PAD   = 16

def prepare_images():
    fond = cv2.imread(FOND_PATH, cv2.IMREAD_GRAYSCALE)
    if fond is None:
        raise FileNotFoundError(f"Impossible de lire {FOND_PATH}")
    fond_small = cv2.resize(fond, None, fx=SCALE_COARSE, fy=SCALE_COARSE, interpolation=cv2.INTER_AREA)

    tiles = []
    for fn in sorted(os.listdir(TILES_DIR)):
        if not fn.lower().endswith(('.png','.jpg','.jpeg')): continue
        img = cv2.imread(os.path.join(TILES_DIR, fn))
        if img is None: continue
        h, w = img.shape[:2]
        inner = img[BORDER_PX:h-BORDER_PX, BORDER_PX:w-BORDER_PX]
        gray  = cv2.cvtColor(inner, cv2.COLOR_BGR2GRAY)
        small = cv2.resize(gray, None, fx=SCALE_COARSE, fy=SCALE_COARSE, interpolation=cv2.INTER_AREA)
        # stocke : nom, gris, petite version, largeur, hauteur
        tiles.append((fn, gray, small, gray.shape[1], gray.shape[0]))
    if not tiles:
        raise RuntimeError("Aucune tuile trouvée")
    return fond, fond_small, tiles

def match_tile(args):
    name, tile_gray, tile_small, tw, th, fond, fond_small = args
    # coarse
    res = cv2.matchTemplate(fond_small, tile_small, cv2.TM_SQDIFF_NORMED)
    _, _, (xs, ys), _ = cv2.minMaxLoc(res)
    x0 = int(xs / SCALE_COARSE)
    y0 = int(ys / SCALE_COARSE)
    # fenêtre de raffinement
    x1 = max(0, x0 - WINDOW_PAD)
    y1 = max(0, y0 - WINDOW_PAD)
    x2 = min(fond.shape[1] - tw, x0 + WINDOW_PAD)
    y2 = min(fond.shape[0] - th, y0 + WINDOW_PAD)
    roi = fond[y1:y2, x1:x2]
    # refine
    res2 = cv2.matchTemplate(roi, tile_gray, cv2.TM_SQDIFF_NORMED)
    _, _, (dx, dy), _ = cv2.minMaxLoc(res2)
    return (name, x1 + dx, y1 + dy, tw, th)

def sort_grid(matches):
    """
    matches = [(name,x,y,w,h), ...]
    On trie par y, puis on découpe en N lignes de N tuiles,
    enfin on trie chaque ligne par x.
    """
    total = len(matches)
    N = int(sqrt(total))
    if N * N != total:
        raise ValueError(f"{total} tuiles n’est pas un carré parfait")

    # 1) tri global par y
    matches.sort(key=lambda t: t[2])

    # 2) découpage en N lignes de N éléments
    grid = []
    for i in range(N):
        row = matches[i * N : (i + 1) * N]
        # tri de la ligne par x
        row.sort(key=lambda t: t[1])
        grid.append([name for name, *rest in row])

    return grid

def main():
    fond, fond_small, tiles = prepare_images()
    args = [(fn, gray, small, tw, th, fond, fond_small) for fn, gray, small, tw, th in tiles]

    with Pool(min(len(args), cpu_count())) as pool:
        results = pool.map(match_tile, args)

    grid = sort_grid(results)

    # affichage
    print("\nOrdre des tuiles (par ligne) :")
    for row in grid:
        print(' '.join(row))

    N = len(grid)
    print("\nCoins du puzzle :")
    print("HG:", grid[0][0], "  HD:", grid[0][N-1])
    print("BG:", grid[N-1][0], "  BD:", grid[N-1][N-1])

if __name__ == '__main__':
    main()