A script to correct the rotation of pennies using OpenCV

rotate_pennies.py

import argparse
import math
import multiprocessing
import os
import itertools as it

from scipy import ndimage
from skimage.metrics import structural_similarity
import cv2
import numpy as np
import matplotlib.pyplot as plt


class PennyRotator:
    def __init__(self):
        self.face_template = cv2.imread('face_template.png')
        self.face_template = self.normalize(self.face_template)

    def correct_rotation_on_penny(self, i, count, image_path):
        print(f'[{i + 1} / {count}] Correcting rotation for "{image_path}"')
        # Open the image
        original_image = cv2.imread(image_path)
        normalized_image = self.normalize(original_image)

        print(f'[{i + 1}] Starting rotate and compare')

        # Get all angles between 0 and 360 inclusive with an increment of 1 into an array
        angles = np.arange(1, 360, 1)

        best_angle = 0
        best_score = 0

        for angle in angles:
            angle, score = self.rotate_and_compare(normalized_image, angle)
            if score > best_score:
                best_score = score
                best_angle = angle

        print(f'[{i + 1}] Best Angle: {best_angle}, Best Score: {best_score}')
        result_image = self.rotate(original_image, best_angle)
        return result_image, original_image, normalized_image, image_path

    def center(self, image):
        return tuple(np.array(image.shape[1::-1]) / 2)

    def rotate_and_compare(self, image, angle):
        rotated = self.rotate(image, angle)
        score = structural_similarity(self.face_template, rotated, gaussian_weights=True)
        return angle, score

    def rotate(self, image, angle):
        rotation_matrix = cv2.getRotationMatrix2D(self.center(image), angle, 1.0)
        return cv2.warpAffine(image, rotation_matrix, image.shape[1::-1], flags=cv2.INTER_LINEAR)

    def normalize(self, image):
        # Apply CLAHE technique to normalize the image
        image_lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        l_channel, a, b = cv2.split(image_lab)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        l_channel = clahe.apply(l_channel)
        image_lab = cv2.merge((l_channel, a, b))

        image = cv2.cvtColor(cv2.cvtColor(image_lab, cv2.COLOR_LAB2BGR), cv2.COLOR_BGR2GRAY)

        size = 100
        image = cv2.resize(image, (size, size), interpolation=cv2.INTER_AREA)

        # Try Canny approach, but use more and more median filtering while it's too noisy.
        retries_with_filtering = 5
        while True:
            result = cv2.Canny(image, 90, 150, apertureSize=3)
            brightness = cv2.sumElems(result)[0]
            if brightness <= 400_000 or retries_with_filtering <= 0:
                break
            image = ndimage.median_filter(image, 3)
            retries_with_filtering -= 1
        image = result

        # Mask the inner circle
        mask = cv2.circle(np.zeros_like(image), (size // 2, size // 2), int(size // 2.5), (255, 255, 255), -1)
        image = cv2.bitwise_and(image, mask)
        return image


parser = argparse.ArgumentParser()
parser.add_argument("-p", "--plot", help="Show plots", action="store_true")
parser.set_defaults(plot=False)
args = parser.parse_args()


def main():
    rotator = PennyRotator()
    if not os.path.exists('output'):
        os.makedirs('output')

    files = os.listdir('data')
    with multiprocessing.Pool() as pool:
        async_results = []
        for i, file in enumerate(files):
            async_results.append(
                pool.apply_async(rotator.correct_rotation_on_penny, (i, len(files), os.path.join('data', file),)))
        for i, result in enumerate(async_results):
            print(f'[{i + 1} / {len(async_results)}] Processing result ')
            result_image, original_image, normalized_image, image_path = result.get()
            image_path = os.path.basename(image_path)
            cv2.imwrite(os.path.join('output', image_path), result_image)

            if args.plot:
                _, axes = plt.subplots(2, 2, figsize=(8, 8))
                axes = axes.flatten()

                titles = ['Original', 'Result', 'Original Normalized', 'Template Normalized']
                is_gray = [False, False, True, True]
                images = [
                    original_image,
                    result_image,
                    normalized_image,
                    rotator.face_template
                ]

                for axis, image, title, is_gray in zip(axes, images, titles, is_gray):
                    axis.set_title(title)
                    if is_gray:
                        axis.imshow(image, cmap='gray')
                    else:
                        axis.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

                plt.suptitle(image_path)
                plt.show()


if __name__ == '__main__':
    main()

Template image