Count jigsaw puzzle pieces using k-means clustering

count_jigsaw_kmeans.py

# Counts jigsaw puzzle pieces
# See https://www.kleemans.ch/counting-jigsaw-puzzle-pieces
# Uses opencv-python==4.11.0.86 and numpy==2.2.2

import statistics

import cv2
import numpy as np


def kmeans_color_quantization(image, clusters=8, rounds=1):
    h, w = image.shape[:2]
    samples = np.zeros([h * w, 3], dtype=np.float32)
    count = 0
    for x in range(h):
        for y in range(w):
            samples[count] = image[x][y]
            count += 1

    compactness, labels, centers = cv2.kmeans(samples, clusters, None,
                                              (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10000, 0.0001),
                                              rounds, cv2.KMEANS_RANDOM_CENTERS)
    centers = np.uint8(centers)
    res = centers[labels.flatten()]
    return res.reshape(image.shape)


# Load image and perform kmeans
filename = 'test'
image = cv2.imread(f'{filename}.jpg')
image_orig = image

# k-means color quantization
kmeans = kmeans_color_quantization(image, clusters=2)
result = kmeans.copy()
cv2.imwrite(f'{filename}_kmeans.png', kmeans)

# Find colors
print('Finding colors...')
image = kmeans

clusters = []
print('Clustering...')
# Convert to grey, threshold
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# Threshold
_, image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

# Invert if needed
print('First pixel:', image[0][0])
if image[0][0] == 255:
    print('Inverting image!')
    image = cv2.bitwise_not(image)
cv2.imwrite(f'{filename}_kmeans_threshold.png', image)

# Laplace
img_laplace2 = cv2.Laplacian(image, cv2.CV_8U, ksize=3)
cv2.imwrite(f'{filename}_kmeans_laplace.png', img_laplace2)

# Find contours
contours, hierarchy = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

# Drawing contours
img_contours = np.zeros(image_orig.shape)
cv2.drawContours(img_contours, contours, -1, (0, 255, 0), 1)
cv2.imwrite(f'{filename}_kmeans_contours.png', img_contours)

# Final piece counting
all_counts = sorted([len(c) for c in contours])
print('Total initial contours:', len(all_counts))
print('All pixel counts:', all_counts)
median = statistics.median(all_counts)
print('Median:', median)

contour_count = 0
for count in all_counts:
    for i in range(10):
        if count < (i + 0.5) * median:
            contour_count += i
            break
print('Final contour count:', contour_count)